Treap模板-CSDN博客

本文链接：https://blog.csdn.net/dong_pope/article/details/103946484
#include <iostream>
#include <algorithm>

const int INF = 0x7fffffff;
const int MAXN = 100005;
using namespace std;

int root = 0;
struct {
    int lef, rgh;
    int val, key;
    int cnt, sub;
} treap[MAXN];
int t = 0;

void Updata(int p) {
    treap[p].sub = treap[treap[p].lef].sub + treap[treap[p].rgh].sub + treap[p].cnt;
}

void zig(int &p) {
    int tmp = treap[p].lef;
    treap[p].lef = treap[tmp].rgh;
    treap[tmp].rgh = p;
    p = tmp;
    Updata(treap[p].rgh);
    Updata(p);
}

void zag(int &p) {
    int tmp = treap[p].rgh;
    treap[p].rgh = treap[tmp].lef;
    treap[tmp].lef = p;
    p = tmp;
    Updata(treap[p].lef);
    Updata(p);
}

void Insert(int &p, int x) {
    if (p == 0) {
        treap[++t] = {0, 0, x, rand(), 1, 1};
        p = t;
        return;
    }
    if (x == treap[p].val) {
        treap[p].cnt++;
        Updata(p);
        return;
    }
    if (x < treap[p].val) {
        Insert(treap[p].lef, x);
        if (treap[p].key < treap[treap[p].lef].key)zig(p);
    }
    if (x > treap[p].val) {
        Insert(treap[p].rgh, x);
        if (treap[p].key < treap[treap[p].rgh].key)zag(p);
    }
    Updata(p);
}

void Erase(int &p, int x) {
    if (p == 0)return;
    if (x == treap[p].val) {
        if (treap[p].cnt > 1) {
            treap[p].cnt--;
            Updata(p);
            return;
        }
        if (treap[p].lef || treap[p].rgh) {
            if (treap[p].rgh == 0 || treap[treap[p].lef].key > treap[treap[p].rgh].key) {
                zig(p);
                Erase(treap[p].rgh, x);
            } else {
                zag(p);
                Erase(treap[p].lef, x);
            }
        } else p = 0;
    } else if (x < treap[p].val)Erase(treap[p].lef, x);
    else if (x > treap[p].val)Erase(treap[p].rgh, x);
    Updata(p);
}

int Getval(int p, int x) {
    if (p == 0)return 0;
    if (x == treap[p].val)return treap[treap[p].lef].sub + 1;
    if (x < treap[p].val)return Getval(treap[p].lef, x);
    return Getval(treap[p].rgh, x) + treap[treap[p].lef].sub + treap[p].cnt;
}

int Getrank(int p, int x) {
    if (p == 0)return INF;
    if (treap[treap[p].lef].sub >= x)return Getrank(treap[p].lef, x);
    if (treap[treap[p].lef].sub + treap[p].cnt >= x)return treap[p].val;
    return Getrank(treap[p].rgh, x - treap[treap[p].lef].sub - treap[p].cnt);
}

int Getpre(int p, int x) {
    if (p == 0)return -INF;
    if (x <= treap[p].val)return Getpre(treap[p].lef, x);
    return max(treap[p].val, Getpre(treap[p].rgh, x));
}

int Getnext(int p, int x) {
    if (p == 0)return INF;
    if (x >= treap[p].val)return Getnext(treap[p].rgh, x);
    return min(treap[p].val, Getnext(treap[p].lef, x));
}

signed main() {
    ios::sync_with_stdio(0);
    cin.tie(NULL);
    int n;
    cin >> n;
    for (int i = 1; i <= n; i++) {
        int op, x;
        cin >> op >> x;
        if (op == 1)Insert(root, x);
        if (op == 2)Erase(root, x);
        if (op == 3)cout << Getval(root, x) << endl;
        if (op == 4)cout << Getrank(root, x) << endl;
        if (op == 5)cout << Getpre(root, x) << endl;
        if (op == 6)cout << Getnext(root, x) << endl;
    }
    return 0;
}