253 普通平衡树

您需要写一种数据结构（可参考题目标题），来维护一些数，其中需要提供以下操作：

1. 插入数值 $x$。
2. 删除数值 $x$(若有多个相同的数，应只删除一个)。
3. 查询数值 $x$ 的排名(若有多个相同的数，应输出最小的排名)。
4. 查询排名为 $x$ 的数值。
5. 求数值 $x$ 的前驱(前驱定义为小于 $x$ 的最大的数)。
6. 求数值 $x$ 的后继(后继定义为大于 $x$ 的最小的数)。

注意：  数据保证查询的结果一定存在。

输入格式

第一行为 $n$，表示操作的个数。

接下来 $n$ 行每行有两个数 $opt$ 和 $x$，$opt$ 表示操作的序号($1\le opt\le 6$)。

输出格式

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

数据范围

$1\le n\le 100000$,所有数均在 $-10^7$ 到 $10^7$ 内。

输入样例：

8
1 10
1 20
1 30
3 20
4 2
2 10
5 25
6 -1

输出样例：

2
20
20
20

思路

用 Treap 就可以实现。 定义数据结构：

struct Node{
	int l, r; // 左右儿子
	int key, val; // key 是节点的值, val是用来确保BST为随机数的堆
	int size, cnt; // size 用来实现排名
}

因为只有随机的 BST, 其高度期望才是 $\log n$, 故我们需要确保其随机, 定义一个 get_node 函数来创建节点：

int get_node(int key)
{
	tr[++idx].key = key;
	tr[idx].value = rand(); // 随机的关键, 因为Treap的堆性质需要满足 u.val > l.val, r.val, 故这样可以实现数据结构打乱, 避免出现一条链的情况
	tr[idx].size = tr[idx].cnt = 1;
	return idx;
}

初始化时, 只需要创建两个边界节点 INF 和 -INF 即可。因为需要实现排名的查询即整个区间中第 k 个数, 故对于点 u 我们需要知道它左边有多少个点, 才能快速知道点 u 的排名。 故定义 pushup 操作来用子节点的 size 更新父节点的 size：

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

左旋/右旋操作。左旋时让当前节点$p$的左子树$tr[p].l$替换掉当前节点, 然后令 $tr[p].l$ 的右子树 $tr[tr[p].l].r$ 作为 $p$ 节点的新左子树。右旋同理。

还需要实现插入一个数的操作：

void insert(int &p, int key)
{
	if(!p) p = get_node(key); // 不存在这个节点就新建一个
	else if(tr[p].key == k) tr[p].cnt++; // 若存在重复值就直接累加个数
	else if(tr[p].key > k)
	{
		insert(tr[p].l, key);
		if(tr[tr[p].l].val > tr[p].val) zig(p); // 保证是一个堆
	}
	else
	{
		insert(tr[p].r, key);
		if(tr[tr[p].r].val > tr[p].val) zag(p);
	}
	pushup(p);
}

可以发现一个规律, 每个操作基本都是先判断当前节点是否存在, 接着判断是否和条件相等, 然后再判断 > 情况, 最后处理 < 情况。

删除操作：

int remove(int &p, int key)
{
	if(!p) return 0;
	if(tr[p].key == key)
	{
		if(tr[p].cnt > 1) tr[p].cnt--;
		else if(tr[p].l || tr[p.r]) // 若存在子树, 就不是叶子结点
		{
			if(!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val) // 若没有右子树或者左子树的 val 大于右子树的 val, 则先右旋再删除旋转后的右子树
			{
				zig(p);
				remove(tr[p].r);
			}
			else
			{
				zag(p);
				remove(tr[p].l);
			}
		}
		else p = 0; // 是叶子节点, 就删去
	}
	else if(tr[p].key > key) remove(tr[p].l, key);
	else remove(tr[p].r.key);
	pushup(p);
}

至于根据值查rank, 也是类似地找到那个点, 若当前点就是, 那么返回 $tr[tr[p].l].size+1$。若在左子树里就递归找左子树, 在右子树的话需要将左子树的rank+当前节点的个数cnt 之和加上。

根据rank查值的话, 先判断左子树中的点个数是否 >= rank, 若满足则说明答案还在左子树中, 若不满足则有两种情况, 一种是当前节点, 另一种是在右子树中。若是当前节点, 那么左子树的节点个数+当前节点的cnt 一定 >= rank, 若不满足则肯定在右子树中, 递归时需要让rank减去左子树点个数和当前点cnt值。

找前驱则是先递归找左子树, 直到当前节点的值 < key, 则搜右子树, 显然如果搜了一次右子树, 会一直去搜右子树, 那么就实现了查找比key小的最大值, 找后驱也是同理。

代码

#include <iostream>
#include <cstring>
#include <algorithm>
#include <string>
#include <cmath>
#include <ctime>
//#include <rand>
using namespace std;
const int N = 1e5 + 10, INF = 0x3f3f3f3f;
int idx;
int root;
struct Node{
    int l, r;
    int key, val;
    int size, cnt;
}tr[N];
int n;
void pushup(int p)
{
    tr[p].size = tr[tr[p].l].size + tr[tr[p].r].size + tr[p].cnt;
}
 
int get_node(int key)
{
    tr[++idx].key = key;
    tr[idx].val = rand();
    tr[idx].cnt = tr[idx].size = 1;
    return idx;
}
 
void build()
{
    get_node(-INF), get_node(INF);
    root = 1, tr[root].r = 2;
    pushup(root);
}
 
void zig(int &p) // 右旋
{
    int q = tr[p].l;
    tr[p].l = tr[q].r;
    tr[q].r = p;
    p = q;
    pushup(tr[p].r), pushup(p);
}
 
void zag(int &p) // 左旋
{
    int q = tr[p].r;
    tr[p].r = tr[q].l;
    tr[q].l = p;
    p = q;
    pushup(tr[p].l), pushup(p);
}
 
void insert(int &p, int key)
{
    if(!p) p = get_node(key);
    else if(tr[p].key == key) tr[p].cnt++;
    else if(tr[p].key > key)
    {
        insert(tr[p].l, key);
        if(tr[tr[p].l].val > tr[p].val) zig(p);
    }
    else {
        insert(tr[p].r, key);
        if(tr[tr[p].r].val > tr[p].val) zag(p);
    }
    pushup(p);
}
 
int remove(int &p, int key)
{
    if(!p) return 0;
    if(tr[p].key == key)
    {
        if(tr[p].cnt > 1) tr[p].cnt --;
        else if(tr[p].l || tr[p].r)
        {
            if(!tr[p].r || tr[tr[p].l].val > tr[tr[p].r].val)
            {
                zig(p);
                remove(tr[p].r,key);
            }
            else
            {
                zag(p);
                remove(tr[p].l, key);
            }
        }
        else p = 0;
    }
    else if(tr[p].key > key) remove(tr[p].l, key);
    else remove(tr[p].r, key);
    pushup(p);
}
 
int get_rank_by_val(int p, int key)
{
    if(!p) return 0;
    if(tr[p].key == key) return tr[tr[p].l].size + 1;
    if(tr[p].key > key) return get_rank_by_val(tr[p].l, key);
    return tr[tr[p].l].size + tr[p].cnt + get_rank_by_val(tr[p].r, key);
}
 
int get_val_by_rank(int p, int size)
{
    if(!p) return INF;
    else if(tr[tr[p].l].size >= size) return get_val_by_rank(tr[p].l, size);
    else if(tr[tr[p].l].size + tr[p].cnt >= size) return tr[p].key;
    return get_val_by_rank(tr[p].r, size - tr[tr[p].l].size - tr[p].cnt);
}
 
int get_prev(int &p, int key)
{
    if(!p) return -INF;
    if(tr[p].key >= key) return get_prev(tr[p].l, key);
    return max(tr[p].key, get_prev(tr[p].r, key));
}
 
int get_next(int &p, int key)
{
    if(!p) return INF;
    if(tr[p].key <= key) return get_next(tr[p].r, key);
    return min(tr[p].key, get_next(tr[p].l, key));
}
 
int main()
{
    cin >> n;
    while(n--)
    {
        int op, x;
        cin >> op >> x;
        if(op == 1) insert(root, x);
        else if(op == 2) remove(root, x);
        else if(op == 3) cout << get_rank_by_val(root, x) << endl;
        else if(op == 4) cout << get_val_by_rank(root, x) << endl;
        else if(op == 5) cout << get_prev(root, x) << endl;
        else cout << get_next(root, x) << endl;
    }
    return 0;
}