test/data_structure/ordered_set.0.test.cpp

• View this file on GitHub
• Last update: 2026-04-09 22:09:49+08:00
• Problem: https://judge.yosupo.jp/problem/ordered_set

Depends on

• node_pool.hpp
• rng.hpp
• treap_node_base.hpp

Code

// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/ordered_set

#include "node_pool.hpp"
#include "treap_node_base.hpp"
#include <iostream>

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    struct TreapNode : TreapNodeBase<TreapNode> {
        int Val;
        explicit TreapNode(int Val) : Val(Val) {}
        bool operator<(const TreapNode &other) const { return Val < other.Val; }
    };
    int n, q;
    std::cin >> n >> q;
    DynamicSizeNodePool<TreapNode> pool;
    TreapNode *root = nullptr;
    for (int i = 0; i < n; ++i) {
        int v;
        std::cin >> v;
        const TreapNode t(v);
        TreapNode *found = TreapNode::find(root, &t);
        if (!found) root = TreapNode::insert(root, pool.make(v));
    }
    while (q--) {
        int cmd, x;
        std::cin >> cmd >> x;
        const TreapNode t(x);
        switch (cmd) {
        case 0: {
            TreapNode *found = TreapNode::find(root, &t);
            if (!found) root = TreapNode::insert(root, pool.make(t.Val));
            break;
        }
        case 1: {
            TreapNode *found = TreapNode::find(root, &t);
            if (found) {
                auto [a, b, c] = TreapNode::split3(root, &t);
                TreapNode *d   = b;
                b              = TreapNode::join(b->left(), b->right());
                pool.retrieve(d);
                root = TreapNode::join(a, b, c);
            }
            break;
        }
        case 2: {
            if (root && root->size() >= x) {
                std::cout << root->select(x - 1)->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        case 3: {
            std::cout << TreapNode::count_less_equal(root, &t) << '\n';
            break;
        }
        case 4: {
            if (TreapNode *found = TreapNode::find(root, &t)) {
                std::cout << found->Val << '\n';
            } else if (TreapNode *pred = TreapNode::predecessor(root, &t)) {
                std::cout << pred->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        case 5: {
            if (TreapNode *found = TreapNode::find(root, &t)) {
                std::cout << found->Val << '\n';
            } else if (TreapNode *succ = TreapNode::successor(root, &t)) {
                std::cout << succ->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        default: break;
        }
    }
    return 0;
}

#line 1 "test/data_structure/ordered_set.0.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/ordered_set

#line 2 "node_pool.hpp"

#include <list>
#include <memory>
#include <utility>
#include <vector>

template<typename NodeT> class FixedSizeNodePool {
    std::vector<NodeT> pool;

public:
    explicit FixedSizeNodePool(int n) : pool(n) {}
    NodeT *at(int ind) { return pool.data() + ind; }
    int id(NodeT *a) const { return a - pool.data(); }
    auto get_func() {
        return std::make_pair([this](int ind) { return at(ind); },
                              [this](NodeT *a) { return id(a); });
    }
};

template<typename NodeT> class DynamicSizeNodePool {
    struct Wrapped : public NodeT {
        using NodeT::NodeT;
        typename std::list<Wrapped>::iterator i;
    };
    std::list<Wrapped> used_list;
    std::list<Wrapped> free_list;

public:
    template<typename... Args> NodeT *make(Args &&...arg) {
        if (free_list.empty()) {
            auto &&node = used_list.emplace_back(std::forward<Args>(arg)...);
            node.i      = std::prev(used_list.end());
            return std::addressof(node);
        }
        used_list.splice(used_list.end(), free_list, std::prev(free_list.end()));
        auto &&node = used_list.back();
        node.~NodeT(); // i remains unchanged
        new (std::addressof(node)) NodeT(std::forward<Args>(arg)...);
        return std::addressof(node);
    }
    // this is lazy, if sth. relies on the order of dtor, do NOT use
    void retrieve(NodeT *node) {
        free_list.splice(free_list.end(), used_list, ((Wrapped *)node)->i);
    }
};
#line 2 "treap_node_base.hpp"

#line 2 "rng.hpp"

#include <cstdint>
#include <limits>

// see: https://prng.di.unimi.it/xoshiro256starstar.c
// original license CC0 1.0
class xoshiro256starstar {
    using u64 = std::uint64_t;

    static inline u64 rotl(const u64 x, int k) { return (x << k) | (x >> (64 - k)); }

    u64 s_[4];

    u64 next() {
        const u64 res = rotl(s_[1] * 5, 7) * 9;
        const u64 t   = s_[1] << 17;
        s_[2] ^= s_[0];
        s_[3] ^= s_[1];
        s_[1] ^= s_[2];
        s_[0] ^= s_[3];
        s_[2] ^= t;
        s_[3] = rotl(s_[3], 45);
        return res;
    }

public:
    // see: https://prng.di.unimi.it/splitmix64.c
    // original license CC0 1.0
    explicit xoshiro256starstar(u64 seed) {
        for (int i = 0; i < 4; ++i) {
            u64 z = (seed += 0x9e3779b97f4a7c15);
            z     = (z ^ (z >> 30)) * 0xbf58476d1ce4e5b9;
            z     = (z ^ (z >> 27)) * 0x94d049bb133111eb;
            s_[i] = z ^ (z >> 31);
        }
    }
    // see: https://en.cppreference.com/w/cpp/named_req/UniformRandomBitGenerator
    using result_type = u64;
    static constexpr u64 min() { return std::numeric_limits<u64>::min(); }
    static constexpr u64 max() { return std::numeric_limits<u64>::max(); }
    u64 operator()() { return next(); }
};
#line 4 "treap_node_base.hpp"
#include <array>
#include <random>
#line 7 "treap_node_base.hpp"

template<typename FlippableTreapNodeT> class FlippableTreapNodeBase;

template<typename TreapNodeT> class TreapNodeBase {
    friend class FlippableTreapNodeBase<TreapNodeT>;

    TreapNodeBase *L;
    TreapNodeBase *R;
    const int Rank;
    int Size;

    TreapNodeT &underlying() { return (TreapNodeT &)*this; }
    const TreapNodeT &underlying() const { return (const TreapNodeT &)*this; }

    static inline xoshiro256starstar gen{std::random_device{}()};
    static inline std::uniform_int_distribution<int> dis{0, 998244353 - 1};

    // CRTP reimplement
    void do_propagate() {}
    void do_update() {}

    // base_propagate() is called to propagate the update information to child(ren).
    // There is no need to update the information combined from child(ren)
    // which should be done in base_update().
    void base_propagate() { underlying().do_propagate(); }
    // base_update() is called to update the information combined from child(ren).
    void base_update() {
        Size = 1;
        if (L) Size += L->Size;
        if (R) Size += R->Size;
        underlying().do_update();
    }

    static TreapNodeBase *base_join(TreapNodeBase *a, TreapNodeBase *b) {
        if (a == nullptr) {
            if (b) b->propagate(), b->update();
            return b;
        }
        if (b == nullptr) {
            if (a) a->propagate(), a->update();
            return a;
        }
        a->propagate(), b->propagate();
        // MAX HEAP
        if (a->rank() < b->rank()) return b->L = base_join(a, b->L), b->update(), b;
        return a->R = base_join(a->R, b), a->update(), a;
    }

    static std::array<TreapNodeBase *, 2> base_split(TreapNodeBase *a, int k) {
        if (a == nullptr) return {nullptr, nullptr};
        a->propagate();
        if (k == 0) return {nullptr, a};
        if (k == a->Size) return {a, nullptr};
        const int leftsize = a->L != nullptr ? a->L->Size : 0;
        if (leftsize < k) {
            auto [b, c] = base_split(a->R, k - leftsize - 1);
            a->R        = b, a->update();
            return {a, c};
        }
        auto [b, c] = base_split(a->L, k);
        a->L        = c, a->update();
        return {b, a};
    }

    TreapNodeBase *base_rotate_left() {
        /*    x              b
             / \            / \
            a   b    =>    x   d
               / \        / \
              c   d      a   c    */
        TreapNodeBase *b = R;
        return b->propagate(), R = b->L, b->L = this, update(), b->update(), b;
    }
    TreapNodeBase *base_rotate_right() {
        /*    x              a
             / \            / \
            a   b    =>    c   x
           / \                / \
          c   d              d   b */
        TreapNodeBase *a = L;
        return a->propagate(), L = a->R, a->R = this, update(), a->update(), a;
    }

protected:
    TreapNodeBase() : L(), R(), Rank(dis(gen)), Size(1) {}

public:
    int size() const { return Size; }
    int rank() const { return Rank; }
    TreapNodeT *left() const { return (TreapNodeT *)L; }
    TreapNodeT *right() const { return (TreapNodeT *)R; }
    void update() { base_update(); }
    void propagate() { underlying().base_propagate(); }
    static TreapNodeT *insert(TreapNodeT *root, TreapNodeT *t) {
        if (root == nullptr) return t;
        root->propagate();
        if (std::as_const(*t) < std::as_const(*root)) {
            root->L = insert(root->left(), t), root->update();
            if (root->rank() < root->left()->rank()) return (TreapNodeT *)root->base_rotate_right();
        } else {
            root->R = insert(root->right(), t), root->update();
            if (root->rank() < root->right()->rank()) return (TreapNodeT *)root->base_rotate_left();
        }
        return root;
    }
    static int count_less_than(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return 0;
        root->propagate();
        if (std::as_const(*root) < *t)
            return (root->left() ? root->left()->size() : 0) + 1 +
                   count_less_than(root->right(), t);
        return count_less_than(root->left(), t);
    }
    static int count_less_equal(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return 0;
        root->propagate();
        if (*t < std::as_const(*root)) return count_less_equal(root->left(), t);
        return (root->left() ? root->left()->size() : 0) + 1 + count_less_equal(root->right(), t);
    }
    // [<, >=]
    static std::array<TreapNodeT *, 2> split_less_than(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return {nullptr, nullptr};
        root->propagate();
        if (std::as_const(*root) < *t) {
            auto [a, b] = split_less_than(root->right(), t);
            root->R     = a, root->update();
            return {root, b};
        }
        auto [a, b] = split_less_than(root->left(), t);
        root->L     = b, root->update();
        return {a, root};
    }
    static int count_greater_than(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return 0;
        root->propagate();
        if (*t < std::as_const(*root))
            return (root->right() ? root->right()->size() : 0) + 1 +
                   count_greater_than(root->left(), t);
        return count_greater_than(root->right(), t);
    }
    static int count_greater_equal(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return 0;
        root->propagate();
        if (std::as_const(*root) < *t) return count_greater_equal(root->right(), t);
        return (root->right() ? root->right()->size() : 0) + 1 +
               count_greater_equal(root->left(), t);
    }
    // [<=, >]
    static std::array<TreapNodeT *, 2> split_less_equal(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return {nullptr, nullptr};
        root->propagate();
        if (*t < std::as_const(*root)) {
            auto [a, b] = split_less_equal(root->left(), t);
            root->L     = b, root->update();
            return {a, root};
        }
        auto [a, b] = split_less_equal(root->right(), t);
        root->R     = a, root->update();
        return {root, b};
    }
    static int count(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return 0;
        root->propagate();
        if (*t < std::as_const(*root)) return count(root->left(), t);
        if (std::as_const(*root) < *t) return count(root->right(), t);
        int res = 1;
        if (root->left()) res += root->left()->size() - count_less_than(root->left(), t);
        if (root->right()) res += root->right()->size() - count_greater_than(root->right(), t);
        return res;
    }
    static bool contains(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return false;
        root->propagate();
        if (*t < std::as_const(*root)) return contains(root->left(), t);
        if (std::as_const(*root) < *t) return contains(root->right(), t);
        return true;
    }
    static std::array<int, 3> count3(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return {0, 0, 0};
        root->propagate();
        if (*t < std::as_const(*root)) {
            auto [a, b, c] = count3(root->left(), t);
            return {a, b, c + 1 + (root->right() ? root->right()->size() : 0)};
        }
        if (std::as_const(*root) < *t) {
            auto [a, b, c] = count3(root->right(), t);
            return {a + 1 + (root->left() ? root->left()->size() : 0), b, c};
        }
        const int a = count_less_than(root->left(), t);
        const int c = count_greater_than(root->right(), t);
        return {a, root->size() - a - c, c};
    }
    static std::array<TreapNodeT *, 3> split3(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return {nullptr, nullptr, nullptr};
        root->propagate();
        if (*t < std::as_const(*root)) {
            auto [a, b, c] = split3(root->left(), t);
            root->L        = c, root->update();
            return {a, b, root};
        }
        if (std::as_const(*root) < *t) {
            auto [a, b, c] = split3(root->right(), t);
            root->R        = a, root->update();
            return {root, b, c};
        }
        auto [a, b] = split_less_than(root->left(), t);
        auto [c, d] = split_less_equal(root->right(), t);
        root->L = b, root->R = c, root->update();
        return {a, root, d};
    }
    static TreapNodeT *predecessor(TreapNodeT *root, const TreapNodeT *t) {
        TreapNodeT *res = nullptr;
        while (root) {
            root->propagate();
            if (std::as_const(*root) < *t) {
                res = root, root = root->right();
            } else {
                root = root->left();
            }
        }
        return res;
    }
    static TreapNodeT *successor(TreapNodeT *root, const TreapNodeT *t) {
        TreapNodeT *res = nullptr;
        while (root) {
            root->propagate();
            if (*t < std::as_const(*root)) {
                res = root, root = root->left();
            } else {
                root = root->right();
            }
        }
        return res;
    }

    template<typename... Nodes> static TreapNodeT *join(Nodes... root) {
        struct Helper {
            TreapNodeBase *Val;
            Helper &operator|(TreapNodeBase *A) {
                Val = base_join(Val, A);
                return *this;
            }
        } nil{nullptr};
        return (TreapNodeT *)(nil | ... | root).Val;
    }
    template<typename... Parts>
    static std::array<TreapNodeT *, sizeof...(Parts) + 1> split(TreapNodeT *root, Parts... part) {
        std::array<TreapNodeT *, sizeof...(Parts) + 1> res;
        res[0]    = root;
        int index = 0;
        (
            [&](int s) {
                auto [l, r]  = base_split(res[index], s);
                res[index]   = (TreapNodeT *)l;
                res[++index] = (TreapNodeT *)r;
            }(part),
            ...);
        return res;
    }

    static TreapNodeT *find(TreapNodeT *root, const TreapNodeT *t) {
        if (root == nullptr) return nullptr;
        root->propagate();
        if (std::as_const(*root) < *t) return find(root->right(), t);
        if (*t < std::as_const(*root)) return find(root->left(), t);
        return root;
    }

    TreapNodeT *select(int k) {
        propagate();
        const int leftsize = left() ? left()->size() : 0;
        if (k == leftsize) return (TreapNodeT *)this;
        if (k < leftsize) return left()->select(k);
        return right()->select(k - leftsize - 1);
    }
};

template<typename FlippableTreapNodeT> class FlippableTreapNodeBase
    : public TreapNodeBase<FlippableTreapNodeT> {
    friend class TreapNodeBase<FlippableTreapNodeT>;

    bool NeedFlip;

    FlippableTreapNodeT &underlying() { return (FlippableTreapNodeT &)*this; }
    const FlippableTreapNodeT &underlying() const { return (const FlippableTreapNodeT &)*this; }

    // CRTP reimplement
    void do_flip() {}

    void base_flip() {
        NeedFlip = !NeedFlip;
        std::swap(this->L, this->R);
        underlying().do_flip();
    }
    void base_propagate() {
        underlying().do_propagate();
        if (NeedFlip) {
            NeedFlip = false;
            if (this->left()) this->left()->underlying().base_flip();
            if (this->right()) this->right()->underlying().base_flip();
        }
    }

protected:
    FlippableTreapNodeBase() : NeedFlip() {}

public:
    void flip() { base_flip(); }
};
#line 5 "test/data_structure/ordered_set.0.test.cpp"
#include <iostream>

int main() {
    std::ios::sync_with_stdio(false);
    std::cin.tie(nullptr);
    struct TreapNode : TreapNodeBase<TreapNode> {
        int Val;
        explicit TreapNode(int Val) : Val(Val) {}
        bool operator<(const TreapNode &other) const { return Val < other.Val; }
    };
    int n, q;
    std::cin >> n >> q;
    DynamicSizeNodePool<TreapNode> pool;
    TreapNode *root = nullptr;
    for (int i = 0; i < n; ++i) {
        int v;
        std::cin >> v;
        const TreapNode t(v);
        TreapNode *found = TreapNode::find(root, &t);
        if (!found) root = TreapNode::insert(root, pool.make(v));
    }
    while (q--) {
        int cmd, x;
        std::cin >> cmd >> x;
        const TreapNode t(x);
        switch (cmd) {
        case 0: {
            TreapNode *found = TreapNode::find(root, &t);
            if (!found) root = TreapNode::insert(root, pool.make(t.Val));
            break;
        }
        case 1: {
            TreapNode *found = TreapNode::find(root, &t);
            if (found) {
                auto [a, b, c] = TreapNode::split3(root, &t);
                TreapNode *d   = b;
                b              = TreapNode::join(b->left(), b->right());
                pool.retrieve(d);
                root = TreapNode::join(a, b, c);
            }
            break;
        }
        case 2: {
            if (root && root->size() >= x) {
                std::cout << root->select(x - 1)->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        case 3: {
            std::cout << TreapNode::count_less_equal(root, &t) << '\n';
            break;
        }
        case 4: {
            if (TreapNode *found = TreapNode::find(root, &t)) {
                std::cout << found->Val << '\n';
            } else if (TreapNode *pred = TreapNode::predecessor(root, &t)) {
                std::cout << pred->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        case 5: {
            if (TreapNode *found = TreapNode::find(root, &t)) {
                std::cout << found->Val << '\n';
            } else if (TreapNode *succ = TreapNode::successor(root, &t)) {
                std::cout << succ->Val << '\n';
            } else {
                std::cout << "-1\n";
            }
            break;
        }
        default: break;
        }
    }
    return 0;
}

Test cases

Env	Name	Status	Elapsed	Memory
clang++	example_00	AC	10 ms	24 MB
clang++	example_01	AC	9 ms	26 MB
clang++	max_random_00	AC	154 ms	28 MB
clang++	max_random_01	AC	124 ms	28 MB
clang++	max_random_02	AC	104 ms	28 MB
clang++	max_random_03	AC	77 ms	28 MB
clang++	max_random_04	AC	131 ms	26 MB
clang++	max_random_05	AC	123 ms	26 MB
clang++	max_random_06	AC	120 ms	24 MB
clang++	max_random_07	AC	120 ms	26 MB
clang++	max_random_08	AC	168 ms	24 MB
clang++	max_random_09	AC	138 ms	28 MB
clang++	max_random_10	AC	105 ms	26 MB
clang++	max_random_11	AC	77 ms	26 MB
clang++	max_random_12	AC	126 ms	24 MB
clang++	max_random_13	AC	135 ms	22 MB
clang++	max_random_14	AC	138 ms	26 MB
clang++	max_random_15	AC	144 ms	26 MB
clang++	max_random_16	AC	826 ms	81 MB
clang++	max_random_17	AC	782 ms	87 MB
clang++	max_random_18	AC	1062 ms	103 MB
clang++	max_random_19	AC	1037 ms	85 MB
clang++	max_random_20	AC	880 ms	86 MB
clang++	max_random_21	AC	663 ms	80 MB
clang++	max_random_22	AC	636 ms	82 MB
clang++	max_random_23	AC	623 ms	85 MB
clang++	max_random_24	AC	934 ms	81 MB
clang++	max_random_25	AC	926 ms	89 MB
clang++	max_random_26	AC	1027 ms	102 MB
clang++	max_random_27	AC	1067 ms	83 MB
clang++	max_random_28	AC	862 ms	82 MB
clang++	max_random_29	AC	799 ms	84 MB
clang++	max_random_30	AC	862 ms	85 MB
clang++	max_random_31	AC	849 ms	81 MB
clang++	small_00	AC	79 ms	24 MB
clang++	small_01	AC	82 ms	22 MB
clang++	small_02	AC	84 ms	28 MB

Back to top page