Luo Gu P3369 [template] general balanced tree (Treap)

Topic link: P3369 [template] general balanced tree

The meaning of problems

A data structure configured to meet the six operations given above.

Thinking

Balanced tree

Template question the balance of the tree.

First learn a bit Treap.

Template from the "algorithm contest Step-up Guide" .

Code

#include <bits/stdc++.h>
using namespace std;
const int maxn = 1e5 + 10;
const int inf = 0x3f3f3f3f;

struct Treap {
    int l, r;
    int val, dat;
    int cnt, size;
} a[maxn];
int tot, root;
int n;

int new_node(int val) {
    a[++tot].val = val;
    a[tot].dat = rand();
    a[tot].cnt = a[tot].size = 1;
    return tot;
}

void update(int p) {
    a[p].size = a[a[p].l].size + a[a[p].r].size + a[p].cnt;
}

void build() {
    new_node(-inf), new_node(inf);
    root = 1, a[1].r = 2;
    update(root);
}

int get_rank(int p, int val) {
    if (p == 0) return 0;
    if (val == a[p].val) return a[a[p].l].size + 1;
    if (val < a[p].val) return get_rank(a[p].l, val);
    return get_rank(a[p].r, val) + a[a[p].l].size + a[p].cnt;
}

int get_val(int p, int rank) {
    if (p == 0) return inf;
    if (a[a[p].l].size >= rank) return get_val(a[p].l, rank);
    if (a[a[p].l].size + a[p].cnt >= rank) return a[p].val;
    return get_val(a[p].r, rank - a[a[p].l].size - a[p].cnt);
}

// 右旋
void zig(int &p) {
    int q = a[p].l;
    a[p].l = a[q].r, a[q].r = p, p = q;
    update(a[p].r), update(p);
}

// 左旋
void zag(int &p) {
    int q = a[p].r;
    a[p].r = a[q].l, a[q].l = p, p = q;
    update(a[p].l), update(p);
}

void add(int &p, int val) {
    if (p == 0) {
        p = new_node(val);
        return;
    }
    if (val == a[p].val) {
        a[p].cnt++;
    } else if (val < a[p].val) {
        add(a[p].l, val);
        if (a[p].dat < a[a[p].l].dat) zig(p);
    } else {
        add(a[p].r, val);
        if (a[p].dat < a[a[p].r].dat) zag(p);
    }
    update(p);
}

int get_pre(int val) {
    int ans = 1; // a[1].val==-inf
    int p = root;
    while (p) {
        if (val == a[p].val) {
            if (a[p].l > 0) {
                p = a[p].l;
                while (a[p].r > 0) p = a[p].r;
                ans = p;
            }
            break;
        }
        if (a[p].val < val && a[p].val > a[ans].val) ans = p;
        p = val < a[p].val ? a[p].l : a[p].r;
    }
    return a[ans].val;
}

int get_next(int val) {
    int ans = 2; // a[2].val==inf
    int p = root;
    while (p) {
        if (val == a[p].val) {
            if (a[p].r > 0) {
                p = a[p].r;
                while (a[p].l > 0) p = a[p].l;
                ans = p;
            }
            break;
        }
        if (a[p].val > val && a[p].val < a[ans].val) ans = p;
        p = val < a[p].val ? a[p].l : a[p].r;
    }
    return a[ans].val;
}

// 把非叶子结点旋转至叶子结点后直接删除
void remove(int &p, int val) {
    if (p == 0) return;
    if (val == a[p].val) {
        if (a[p].cnt > 1) {
            a[p].cnt--, update(p);
            return;
        }
        if (a[p].l || a[p].r) { // 不是叶子节点，向下旋转
            if (a[p].r == 0 || a[a[p].l].dat > a[a[p].r].dat)
                zig(p), remove(a[p].r, val);
            else
                zag(p), remove(a[p].l, val);
            update(p);
        }
        else p = 0; // 叶子节点，删除
        return;
    }
    val < a[p].val ? remove(a[p].l, val) : remove(a[p].r, val);
    update(p);
}

int main() {
    build();
    cin >> n;
    while (n--) {
        int op, x;
        cin >> op >> x;

        if(op == 1) {
            add(root, x);
        } else if(op == 2) {
            remove(root, x);
        } else if(op == 3) {
            cout << get_rank(root, x) - 1 << endl;
        } else if(op == 4) {
            cout << get_val(root, x + 1) << endl;
        } else if(op == 5) {
            cout << get_pre(x) << endl;
        } else {
            cout << get_next(x) << endl;
        }
    }
    return 0;
}