CP Notebook

← all categories

Treap Ordered 2f4cc70b

Key-ordered treap: order-statistics multiset via insert/erase/rank/kth/ predecessor/successor, using split-by-key + priority-based merge (no rotations, no parent pointers). T needs operator< and operator==. findKth/predecessor/successor are UB if called out of bounds — guard with size() first. Companion to treap-implicit.h (that one is positional/array-indexed; this one is key-ordered). Multi-page exception to the line budget.

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

content/data-structures/treap-ordered.h

template <typename T>
struct Treap {
    struct Node {
        T val;
        unsigned pri;
        int cnt = 1, sz = 1, l = -1, r = -1;
    };
    vector<Node> pool;
    mt19937 rng{random_device{}()};
    int root = -1;

    int size() { return getSz(root); }
    int getSz(int n) { return n == -1 ? 0 : pool[n].sz; }
    void pull(int n) { pool[n].sz = pool[n].cnt + getSz(pool[n].l) + getSz(pool[n].r); }
    int make(T x) { pool.push_back({x, (unsigned)rng(), 1, 1, -1, -1}); return (int)pool.size() - 1; }

    int merge(int l, int r) {
        if (l == -1) return r;
        if (r == -1) return l;
        if (pool[l].pri > pool[r].pri) { pool[l].r = merge(pool[l].r, r); pull(l); return l; }
        else { pool[r].l = merge(l, pool[r].l); pull(r); return r; }
    }

    void split(int n, T key, int &l, int &r) {  // l: val < key, r: val >= key
        if (n == -1) { l = r = -1; return; }
        if (pool[n].val < key) { split(pool[n].r, key, pool[n].r, r); l = n; }
        else { split(pool[n].l, key, l, pool[n].l); r = n; }
        pull(n);
    }

    int insertRec(int n, int nn) {
        if (n == -1) return nn;
        if (pool[n].val == pool[nn].val) { pool[n].cnt++; pull(n); return n; }
        if (pool[nn].pri > pool[n].pri) {
            split(n, pool[nn].val, pool[nn].l, pool[nn].r);
            pull(nn);
            return nn;
        }
        if (pool[nn].val < pool[n].val) pool[n].l = insertRec(pool[n].l, nn);
        else pool[n].r = insertRec(pool[n].r, nn);
        pull(n);
        return n;
    }
    void insert(T x) { root = insertRec(root, make(x)); }

    int eraseRec(int n, T x) {
        if (n == -1) return -1;
        if (pool[n].val == x) {
            if (pool[n].cnt > 1) { pool[n].cnt--; pull(n); return n; }
            return merge(pool[n].l, pool[n].r);
        }
        if (x < pool[n].val) pool[n].l = eraseRec(pool[n].l, x);
        else pool[n].r = eraseRec(pool[n].r, x);
        pull(n);
        return n;
    }
    void erase(T x) { root = eraseRec(root, x); }

    bool contains(T x) {
        int n = root;
        while (n != -1) {
            if (pool[n].val == x) return true;
            n = pool[n].val < x ? pool[n].r : pool[n].l;
        }
        return false;
    }

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

    T findKth(int k) {  // 1-indexed
        int n = root;
        while (n != -1) {
            int szL = getSz(pool[n].l);
            if (k <= szL) n = pool[n].l;
            else if (k <= szL + pool[n].cnt) return pool[n].val;
            else { k -= szL + pool[n].cnt; n = pool[n].r; }
        }
        return T();
    }

    T predecessor(T x) {
        int n = root, f = -1;
        while (n != -1) {
            if (pool[n].val < x) { f = n; n = pool[n].r; }
            else n = pool[n].l;
        }
        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].l; }
            else n = pool[n].r;
        }
        return pool[f].val;
    }
};