普通平衡树

一本通·提高篇 · 代码 · cpp · 难度 4/5 · 共 4117 字
维护数集，支持：1插入、2删除一个、3求x排
一本通
练习
正文

/*
原题：T1568 「一本通 4.6 练习 3」普通平衡树（Tyvj1728）
题意：维护数集，支持：1 插入、2 删除一个、3 求 x 排名(比 x 小的个数+1)、
4 求第 k 小、5 求 x 前驱(小于 x 的最大)、6 求 x 后继(大于 x 的最小)。
思路：Treap（随机优先级 + split/merge）。排名用 split 求“严格小于 x”的个数 +1；
前驱/后继分别取 split 后左树最大、右树最小；删除只删一个（cnt 计数重复值）。
复杂度：时间 O(n log n)，空间 O(n)。
易错点：删除只删一个相同值；x 不在集合中时 3/5/6 仍按定义返回（前驱/后继可能不存在）；
同值多个时排名与 kth 都要考虑 cnt。split/merge 用下标返回避免 vector 重分配悬垂引用。
*/
#include <bits/stdc++.h>
using namespace std;
struct Treap {
    struct Node {
        int l, r, pri, cnt, sz;
        long long key;
    };
    vector<Node> t;
    int root;
    Treap(){
        t.push_back({0, 0, 0, 0, 0, 0});
        root = 0;
    }
    int rnd(){
        return rand();
    }
    int newNode(long long key){
        t.push_back({0, 0, rnd(), 1, 1, key});
        return (int)t.size() - 1;
    }
    void upd(int x){
        if(x) t[x].sz = t[t[x].l].sz + t[t[x].r].sz + t[x].cnt;
    }
    pair<int, int> split(int x, long long key){
        if(!x) return {0, 0};
        if(t[x].key < key){
            auto p = split(t[x].r, key);
            t[x].r = p.first;
            upd(x);
            return {x, p.second};
        }else{
            auto p = split(t[x].l, key);
            t[x].l = p.second;
            upd(x);
            return {p.first, x};
        }
    }
    int merge(int L, int R){
        if(!L || !R) return L | R;
        if(t[L].pri > t[R].pri){
            t[L].r = merge(t[L].r, R);
            upd(L);
            return L;
        }else{
            t[R].l = merge(L, t[R].l);
            upd(R);
            return R;
        }
    }
    void insert(long long key){
        auto p = split(root, key);
        int L = p.first, R = p.second;
        auto p2 = split(R, key + 1);
        int M = p2.first, R2 = p2.second;
        if(M){
            t[M].cnt++;
            t[M].sz++;
        }else{
            M = newNode(key);
        }
        root = merge(merge(L, M), R2);
    }
    void eraseOne(long long key){
        auto p = split(root, key);
        int L = p.first, R = p.second;
        auto p2 = split(R, key + 1);
        int M = p2.first, R2 = p2.second;
        if(M){
            t[M].cnt--;
            t[M].sz--;
            if(t[M].cnt == 0) M = 0;
        }
        root = merge(merge(L, M), R2);
    }
    int kth(int k){
        int x = root;
        while(x){
            int ls = t[t[x].l].sz;
            if(k <= ls) x = t[x].l;
            else if(k <= ls + t[x].cnt) return (int)t[x].key;
            else {
                k -= ls + t[x].cnt;
                x = t[x].r;
            }
        }
        return -1;
    }
    int rank(long long key){
        auto p = split(root, key);
        int L = p.first, R = p.second;
        int r = t[L].sz;
        root = merge(L, R);
        return r;
    }
    int kthIn(int x, int k){
        while(x){
            int ls = t[t[x].l].sz;
            if(k <= ls) x = t[x].l;
            else if(k <= ls + t[x].cnt) return (int)t[x].key;
            else {
                k -= ls + t[x].cnt;
                x = t[x].r;
            }
        }
        return -1;
    }
    long long pred(long long key){
        auto p = split(root, key);
        int L = p.first, R = p.second;
        long long res = -1;
        if(t[L].sz) res = kthIn(L, t[L].sz);
        root = merge(L, R);
        return res;
    }
    long long succ(long long key){
        auto p = split(root, key + 1);
        int L = p.first, R = p.second;
        long long res = -1;
        if(t[R].sz) res = kthIn(R, 1);
        root = merge(L, R);
        return res;
    }
};
int main(){
    ios::sync_with_stdio(false);
    cin.tie(0);
    int n;
    cin >> n;
    Treap tr;
    while(n--){
        int opt;
        long long x;
        cin >> opt >> x;
        if(opt == 1){
            tr.insert(x);
        } else if(opt == 2){
            tr.eraseOne(x);
        } else if(opt == 3){
            cout << tr.rank(x) + 1 << "\n";
        } else if(opt == 4){
            cout << tr.kth(x) << "\n";
        } else if(opt == 5){
            cout << tr.pred(x) << "\n";
        }else{
            cout << tr.succ(x) << "\n";
        }
    }
    return 0;
}
普通平衡树

正文

一本通·提高篇的其它内容
