test/yosupo/ds/ordered_set.test.cpp

View this file on GitHub
Last update: 2025-09-25 11:45:13+07:00
Problem: https://judge.yosupo.jp/problem/ordered_set

Depends on

Code

#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"

#include <bits/stdc++.h>
using namespace std;

#include "ds/ordered_set.hpp"

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;

    ordered_set<int> s;
    for (int i = 0, a; i < n; i++) cin >> a, s.insert(a);

    while (q--) {
        int t, x;
        cin >> t >> x;
        if (t == 0) {
            s.insert(x);
        } else if (t == 1) {
            s.erase(x);
        } else if (t == 2) {
            cout << s.kth_order(x).value_or(-1) << "\n";
        } else if (t == 3) {
            cout << s.order_of(x) << "\n";
        } else if (t == 4) {
            cout << s.pre_upper_bound(x).value_or(-1) << "\n";
        } else {
            cout << s.lower_bound(x).value_or(-1) << "\n";
        }
    }
}

#line 1 "test/yosupo/ds/ordered_set.test.cpp"
#define PROBLEM "https://judge.yosupo.jp/problem/ordered_set"

#include <bits/stdc++.h>
using namespace std;

#line 2 "ds/ordered_set.hpp"

#line 2 "ds/treap_base.hpp"

#line 2 "ds/bst_base.hpp"

#line 2 "other/st_alloc.hpp"

template<class T> struct st_alloc {
    static inline stack<T> pool;
    void *operator new(size_t) { return pool.emplace(), &pool.top(); }
    void operator delete(void *) {}
};
#line 4 "ds/bst_base.hpp"

namespace bst {
    template<class T> struct node : st_alloc<T> {
        T *l = nullptr, *r = nullptr;
        void pull() { ((T *)this)->_pull(); }
    };
} // namespace bst
#line 4 "ds/treap_base.hpp"

namespace treap {
    template<class T> using key_t = decltype(declval<T>().key);
    template<class T> struct node : bst::node<T> {
        uint64_t pr = rng();

        static uint64_t rng() {
            static mt19937_64 mt(
                chrono::steady_clock::now().time_since_epoch().count());
            return mt();
        }
    };

    template<class T> T *merge(node<T> *_l, node<T> *_r) {
        T *l = (T *)_l, *r = (T *)_r;
        if (!l || !r) return l ? l : r;
        if (l->pr > r->pr) {
            l->r = merge(l->r, r), l->pull();
            return l;
        } else {
            r->l = merge(l, r->l), r->pull();
            return r;
        }
    }
    template<class T>
    pair<T *, T *> split(node<T> *_t, const key_t<T> &k, bool include = true) {
        T *t = (T *)_t;
        if (!t) return {nullptr, nullptr};
        if (k < t->key || (k == t->key && !include)) {
            auto s = split(t->l, k, include);
            t->l = s.second, t->pull();
            return {s.first, t};
        } else {
            auto s = split(t->r, k, include);
            t->r = s.first, t->pull();
            return {t, s.second};
        }
    }
} // namespace treap
#line 4 "ds/ordered_set.hpp"

// NOTE: The copy constructor only copies the pointer. Be careful when doing
// `ordered_set<> s2 = s1;`
namespace oset {
    using treap::merge, treap::split;
    template<class T> struct ordered_set {
        ordered_set() {}

        void insert(const T &k) {
            if (contains(k)) return;
            auto [a, b] = split(root, k, false);
            root = merge(a, merge(new node(k), b));
        }
        void erase(const T &k) {
            if (!contains(k)) return;
            node *a, *b, *c;
            tie(a, b) = split(root, k, false);
            tie(b, c) = split(b, k, true);
            root = merge(a, c);
        }
        optional<T> kth_order(int k) const {
            if (empty() || k > size()) return nullopt;
            int ord = 0;
            for (auto t = root; t;) {
                int w = ord + size(t->l) + 1;
                if (w == k) return t->key;
                t = k < w ? t->l : (ord = w, t->r);
            };
            return nullopt;
        }
        int order_of(const T &k, bool include = true) const {
            if (empty()) return 0;
            int ord = 0;
            for (auto t = root; t;) {
                if (k < t->key || (k == t->key && !include)) {
                    t = t->l;
                } else {
                    ord += size(t->l) + 1, t = t->r;
                }
            }
            return ord;
        }
        optional<T> pre_upper_bound(const T &k) const {
            return kth_order(order_of(k));
        }
        optional<T> lower_bound(const T &k) const {
            return kth_order(order_of(k, false) + 1);
        }

        int size() const { return size(root); }
        bool empty() const { return size() == 0; }
        bool contains(const T &k) const {
            for (auto t = root; t; t = k < t->key ? t->l : t->r)
                if (k == t->key) return true;
            return false;
        }

      private:
        struct node : treap::node<node> {
            using base = treap::node<node>;
            using base::l, base::r;

            T key;
            int sz = 1;
            node() {}
            node(const T k) : key(k) {}
            void _pull() { sz = size(l) + 1 + size(r); }
        };
        node *root = nullptr;
        static int size(node *t) { return t ? t->sz : 0; }
    };
} // namespace oset
using oset::ordered_set;
#line 7 "test/yosupo/ds/ordered_set.test.cpp"

int main() {
    cin.tie(nullptr)->sync_with_stdio(false);
    int n, q;
    cin >> n >> q;

    ordered_set<int> s;
    for (int i = 0, a; i < n; i++) cin >> a, s.insert(a);

    while (q--) {
        int t, x;
        cin >> t >> x;
        if (t == 0) {
            s.insert(x);
        } else if (t == 1) {
            s.erase(x);
        } else if (t == 2) {
            cout << s.kth_order(x).value_or(-1) << "\n";
        } else if (t == 3) {
            cout << s.order_of(x) << "\n";
        } else if (t == 4) {
            cout << s.pre_upper_bound(x).value_or(-1) << "\n";
        } else {
            cout << s.lower_bound(x).value_or(-1) << "\n";
        }
    }
}

Test cases

Env	Name	Status	Elapsed	Memory
g++	example_00	AC	4 ms	4 MB
g++	example_01	AC	4 ms	3 MB
g++	max_random_00	AC	119 ms	5 MB
g++	max_random_01	AC	104 ms	4 MB
g++	max_random_02	AC	79 ms	4 MB
g++	max_random_03	AC	59 ms	4 MB
g++	max_random_04	AC	92 ms	4 MB
g++	max_random_05	AC	92 ms	4 MB
g++	max_random_06	AC	100 ms	4 MB
g++	max_random_07	AC	102 ms	4 MB
g++	max_random_08	AC	124 ms	5 MB
g++	max_random_09	AC	112 ms	4 MB
g++	max_random_10	AC	81 ms	4 MB
g++	max_random_11	AC	59 ms	4 MB
g++	max_random_12	AC	94 ms	4 MB
g++	max_random_13	AC	101 ms	4 MB
g++	max_random_14	AC	121 ms	4 MB
g++	max_random_15	AC	115 ms	4 MB
g++	max_random_16	AC	363 ms	21 MB
g++	max_random_17	AC	376 ms	21 MB
g++	max_random_18	AC	447 ms	26 MB
g++	max_random_19	AC	388 ms	20 MB
g++	max_random_20	AC	415 ms	20 MB
g++	max_random_21	AC	284 ms	20 MB
g++	max_random_22	AC	328 ms	20 MB
g++	max_random_23	AC	308 ms	20 MB
g++	max_random_24	AC	437 ms	21 MB
g++	max_random_25	AC	491 ms	21 MB
g++	max_random_26	AC	502 ms	26 MB
g++	max_random_27	AC	421 ms	20 MB
g++	max_random_28	AC	409 ms	20 MB
g++	max_random_29	AC	411 ms	20 MB
g++	max_random_30	AC	475 ms	20 MB
g++	max_random_31	AC	424 ms	20 MB
g++	small_00	AC	65 ms	5 MB
g++	small_01	AC	67 ms	5 MB
g++	small_02	AC	69 ms	5 MB
clang++	example_00	AC	4 ms	4 MB
clang++	example_01	AC	4 ms	4 MB
clang++	max_random_00	AC	122 ms	5 MB
clang++	max_random_01	AC	109 ms	4 MB
clang++	max_random_02	AC	84 ms	4 MB
clang++	max_random_03	AC	61 ms	4 MB
clang++	max_random_04	AC	101 ms	4 MB
clang++	max_random_05	AC	96 ms	4 MB
clang++	max_random_06	AC	110 ms	4 MB
clang++	max_random_07	AC	106 ms	4 MB
clang++	max_random_08	AC	130 ms	5 MB
clang++	max_random_09	AC	124 ms	4 MB
clang++	max_random_10	AC	85 ms	4 MB
clang++	max_random_11	AC	61 ms	4 MB
clang++	max_random_12	AC	104 ms	4 MB
clang++	max_random_13	AC	100 ms	4 MB
clang++	max_random_14	AC	120 ms	4 MB
clang++	max_random_15	AC	121 ms	4 MB
clang++	max_random_16	AC	462 ms	21 MB
clang++	max_random_17	AC	425 ms	21 MB
clang++	max_random_18	AC	545 ms	26 MB
clang++	max_random_19	AC	491 ms	20 MB
clang++	max_random_20	AC	492 ms	20 MB
clang++	max_random_21	AC	324 ms	20 MB
clang++	max_random_22	AC	359 ms	20 MB
clang++	max_random_23	AC	373 ms	20 MB
clang++	max_random_24	AC	467 ms	21 MB
clang++	max_random_25	AC	477 ms	21 MB
clang++	max_random_26	AC	495 ms	26 MB
clang++	max_random_27	AC	438 ms	20 MB
clang++	max_random_28	AC	461 ms	20 MB
clang++	max_random_29	AC	387 ms	20 MB
clang++	max_random_30	AC	459 ms	20 MB
clang++	max_random_31	AC	463 ms	20 MB
clang++	small_00	AC	66 ms	5 MB
clang++	small_01	AC	68 ms	5 MB
clang++	small_02	AC	68 ms	5 MB