← all categories

Splay Ordered Set 9171fab9

Order-statistics multiset via splay tree: insert/erase/find/rank/kth/ predecessor/successor. T needs operator< and operator==. findKth/predecessor/successor are UB if called out of bounds (empty result) — guard with size() first. Multi-page exception to the line budget.

Time: O(\log n) amortized per operation. tested

content/data-structures/splay-ordered-set.h

template <typename T>
struct Splay {
    struct Node {
        int p = -1, ch[2] = {-1, -1}, sz = 1, cnt = 1;
        T val;
    };
    vector<Node> pool;
    int root = -1;

    int size() { return getSz(root); }
    int getSz(int n) { return n == -1 ? 0 : pool[n].sz; }

    void pull(int n) {
        if (n == -1) return;
        pool[n].sz = pool[n].cnt + getSz(pool[n].ch[0]) + getSz(pool[n].ch[1]);
    }

    int make(T x) { pool.push_back(Node()); pool.back().val = x; return (int)pool.size() - 1; }
    int dir(int n) { return pool[pool[n].p].ch[1] == n; }

    void rotate(int n) {
        int u = pool[n].p, v = pool[u].p, d = dir(n), w = pool[n].ch[d ^ 1];
        pool[u].ch[d] = w;
        if (w != -1) pool[w].p = u;
        pool[n].ch[d ^ 1] = u;
        pool[u].p = n;
        pool[n].p = v;
        if (v != -1) {
            if (pool[v].ch[0] == u) pool[v].ch[0] = n;
            else if (pool[v].ch[1] == u) pool[v].ch[1] = n;
        }
        pull(u); pull(n);
    }

    void splay(int n, int g = -1) {
        while (pool[n].p != g) {
            int u = pool[n].p, v = pool[u].p;
            if (v != g) {
                if (dir(n) == dir(u)) rotate(u);
                else rotate(n);
            }
            rotate(n);
        }
        if (g == -1) root = n;
    }

    int find(T x) {
        int n = root;
        while (n != -1) {
            if (pool[n].val == x) { splay(n); return n; }
            n = pool[n].ch[pool[n].val < x];
        }
        return -1;
    }

    void insert(T x) {
        if (root == -1) { root = make(x); return; }
        int n = root, f = -1;
        while (n != -1) {
            f = n;
            if (pool[n].val == x) { pool[n].cnt++; pull(n); splay(n); return; }
            n = pool[n].ch[pool[n].val < x];
        }
        n = make(x);
        pool[f].ch[pool[f].val < x] = n;
        pool[n].p = f;
        splay(n);
    }

    void erase(T x) {
        int n = find(x);
        if (n == -1) return;
        if (pool[n].cnt > 1) { pool[n].cnt--; pull(n); return; }

        if (pool[n].ch[0] == -1 && pool[n].ch[1] == -1) { root = -1; return; }
        if (pool[n].ch[0] == -1) { root = pool[n].ch[1]; pool[root].p = -1; return; }
        if (pool[n].ch[1] == -1) { root = pool[n].ch[0]; pool[root].p = -1; return; }

        int L = pool[n].ch[0], R = pool[n].ch[1];
        pool[L].p = pool[R].p = -1;
        root = L;
        int nrt = L;
        while (pool[nrt].ch[1] != -1) nrt = pool[nrt].ch[1];
        splay(nrt);
        pool[nrt].ch[1] = R;
        pool[R].p = nrt;
        pull(nrt);
    }

    int findRank(T x) {
        int n = root, res = 1;
        while (n != -1) {
            if (pool[n].val == x) { res += getSz(pool[n].ch[0]); splay(n); return res; }
            if (!(x < pool[n].val)) res += pool[n].cnt + getSz(pool[n].ch[0]);
            n = pool[n].ch[pool[n].val < x];
        }
        return res;
    }

    T findKth(int k) {  // 1-indexed
        int n = root;
        while (n != -1) {
            int szL = getSz(pool[n].ch[0]);
            if (k <= szL) n = pool[n].ch[0];
            else {
                k -= szL;
                if (k <= pool[n].cnt) { splay(n); return pool[n].val; }
                k -= pool[n].cnt; n = pool[n].ch[1];
            }
        }
        return T();  // out of bounds
    }

    T predecessor(T x) {
        int n = root, f = -1;
        while (n != -1) {
            if (pool[n].val < x) { f = n; n = pool[n].ch[1]; }
            else n = pool[n].ch[0];
        }
        if (f != -1) splay(f);
        return pool[f].val;
    }

    T successor(T x) {
        int n = root, f = -1;
        while (n != -1) {
            if (x < pool[n].val) { f = n; n = pool[n].ch[0]; }
            else n = pool[n].ch[1];
        }
        if (f != -1) splay(f);
        return pool[f].val;
    }
};