【洛谷】P3369 【模板】普通平衡树

题目地址：

https://www.luogu.com.cn/problem/P3369

题目描述：
您需要写一种数据结构（可参考题目标题），来维护一些数，其中需要提供以下操作：
插入$x$数
删除$x$数（若有多个相同的数，因只删除一个）
查询$x$数的排名（排名定义为比当前数小的数的个数$+1$）
查询排名为$x$的数
求$x$的前驱（前驱定义为小于$x$，且最大的数）
求$x$的后继（后继定义为大于$x$，且最小的数）

输入格式：
第一行为$n$，表示操作的个数，下面$n$行每行有两个数$\text{opt}$和$x$，$\text{opt}$表示操作的序号（$1\le \text{opt}\le 6$）

输出格式：
对于操作$3,4,5,6$每行输出一个数，表示对应答案

数据范围：
对于$100\%$的数据，$1\le n\le 10^5$，$|x|\le 10^7$

参考https://blog.csdn.net/qq_46105170/article/details/118997891。注意几个小点：
1、insert之后在回溯前需要看一下是否需要旋转，而remove不需要，这是因为remove的时候会将要被删除的节点旋转到下面去，但这并不改变堆性质，所以回溯之前不需要再check一遍，而insert的时候有可能增加的节点违反了堆性质，所以需要check一下是否需要旋转；
2、旋转的时候代码里不需要写pushup当前节点，因为无论在insert还是remove函数里，回溯之前都会pushup当前节点的，不需要在旋转的时候做重复工作。

代码如下：

#include <iostream>
#include <cstring>
using namespace std;

const int N = 1e5 + 10, INF = 1e8;
struct Node {
    int son[2];
    int key, val;
    int cnt, size;
} tr[N];
int n, root, idx;

int get_node(int key) {
    tr[++idx].key = key;
    tr[idx].val = rand();
    tr[idx].cnt = tr[idx].size = 1;
    return idx;
}

void pushup(int p) {
    tr[p].size = tr[tr[p].son[0]].size + tr[tr[p].son[1]].size + tr[p].cnt;
}

// d = 0代表左旋，d = 1代表右旋
void rotate(int& p, int d) {
    int q = tr[p].son[d ^ 1];
    tr[p].son[d ^ 1] = tr[q].son[d], tr[q].son[d] = p, p = q;
	pushup(tr[p].son[d]);
}

void insert(int& p, int key) {
    if (!p) p = get_node(key);
    else if (key < tr[p].key) {
        insert(tr[p].son[0], key);
        if (tr[tr[p].son[0]].val > tr[p].val) rotate(p, 1);
    } else if (key > tr[p].key) {
        insert(tr[p].son[1], key);
        if (tr[tr[p].son[1]].val > tr[p].val) rotate(p, 0);
    } else tr[p].cnt++;

    pushup(p);
}

void remove(int& p, int key) {
    if (!p) return;
    if (key < tr[p].key) remove(tr[p].son[0], key);
    else if (key > tr[p].key) remove(tr[p].son[1], key);
    else if (tr[p].cnt > 1) tr[p].cnt--;
    else if (!tr[p].son[0] && !tr[p].son[1]) p = 0;
    else if (!tr[p].son[1] || (tr[p].son[0] && tr[tr[p].son[0]].val > tr[tr[p].son[1]].val)) {
        rotate(p, 1);
        remove(tr[p].son[1], key);
    } else {
        rotate(p, 0);
        remove(tr[p].son[0], key);
    }

    pushup(p);
}

int get_rank_by_key(int p, int key) {
    if (!p) return 0;
    if (key < tr[p].key) return get_rank_by_key(tr[p].son[0], key);
    else if (key > tr[p].key) return tr[tr[p].son[0]].size + tr[p].cnt + get_rank_by_key(tr[p].son[1], key);
    else return tr[tr[p].son[0]].size + 1;
}

int get_key_by_rank(int p, int rank) {
    if (!p) return INF;
    else if (rank <= tr[tr[p].son[0]].size) return get_key_by_rank(tr[p].son[0], rank);
    else if (rank > tr[tr[p].son[0]].size + tr[p].cnt) return get_key_by_rank(tr[p].son[1], rank - tr[tr[p].son[0]].size - tr[p].cnt);
    else return tr[p].key;
}

int get_prev(int p, int key) {
    if (!p) return -INF;
    if (key <= tr[p].key) return get_prev(tr[p].son[0], key);
    else return max(tr[p].key, get_prev(tr[p].son[1], key));
}

int get_next(int p, int key) {
    if (!p) return INF;
    if (key >= tr[p].key) return get_next(tr[p].son[1], key);
    else return min(tr[p].key, get_next(tr[p].son[0], key));
}

int main() {
    scanf("%d", &n);
    while (n--) {
        int op, x;
        scanf("%d%d", &op, &x);
        if (op == 1) insert(root, x);
        else if (op == 2) remove(root, x);
        else if (op == 3) printf("%d\n", get_rank_by_key(root, x));
        else if (op == 4) printf("%d\n", get_key_by_rank(root, x));
        else if (op == 5) printf("%d\n", get_prev(root, x));
        else if (op == 6) printf("%d\n", get_next(root, x));
    }

    return 0;
}

每个操作期望时间复杂度都是$O(\log n)$，空间$O(n)$，$n$是当前Treap里节点个数。