Ordered Set (ds/ordered_set.hpp)

• View this file on GitHub
• View document part on GitHub
• Last update: 2025-09-25 11:45:13+07:00
• Include: #include "ds/ordered_set.hpp"

Description

A set which can perform order statistic queries.

Operations

• ordered_set<T>(int n)
  • Construct an ordered set
  • Time complexity: $\mathcal O(1)$
• void insert(T k)
  • Inserts k into the set
  • Time complexity: $\mathcal O(\log n)$
• void erase(T k)
  • Erases k from the set
  • Time complexity: $\mathcal O(\log n)$
• optional<T> kth_order(int k)
  • Returns the element which has k elements before it
  • Time complexity: $\mathcal O(\log n)$
• int order_of(T k, bool include = true)
  • Count the number of elements less than (or equal if include == true) to k
  • Time complexity: $\mathcal O(\log n)$
• optional<T> pre_upper_bound(T k)
  • Finds the maximum element less than or equal to k
  • Time complexity: $\mathcal O(\log n)$
• optional<T> lower_bound(T k)
  • Finds the minimum element greater than or equal to k
  • Time complexity: $\mathcal O(\log n)$

Depends on

• ds/bst_base.hpp
• ds/treap_base.hpp
• other/st_alloc.hpp

Verified with

• test/yosupo/ds/ordered_set.test.cpp
• test/yosupo/ds/point_set_range_frequency.test.cpp

Code

#pragma once

#include "../ds/treap_base.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 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;

Back to top page