P3369 普通平衡树模板 treap

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

treap和bst（二叉查找树）的区别就在于 treap 给了每个节点一个随机的权值，在插入的时候，通过该权值使整棵树依然维持一个大根堆/小根堆的形式（使用左旋和右旋方式来完成），在删除的时候，通过将该点旋转到叶子结点后再删除。

代码参考 <算法竞赛进阶指南>

//https://www.luogu.com.cn/problem/P3369
//treap
#include <bits/stdc++.h>
#define ll long long
#define sc scanf
#define pr printf
using namespace std;
const int MAXN = 1e5 + 5;
const ll inf = 1e9 + 7;
struct node
{
	int val;//权值
	int sz;//子树大小
	int ch[2];//左右儿子
	int cnt;//该权值个数
	int data;//随机数（用来平衡二叉树）
}treap[MAXN];
int cnt, root;
int newnode(int val)
{
	treap[cnt].val = val;
	treap[cnt].sz = 1;
	treap[cnt].ch[0] = treap[cnt].ch[1] = 0;
	treap[cnt].cnt = 1;
	treap[cnt].data = rand();
	return cnt++;
}
void up(int rot)
{
	treap[rot].sz = treap[treap[rot].ch[0]].sz + treap[treap[rot].ch[1]].sz + treap[rot].cnt;
}
void build()
{
	cnt = 1;
	newnode(-inf); newnode(inf);
	root = 1;
	treap[1].ch[1] = 2;
	up(1);
}
void rotate(int& rot, int op)
{
	int t = treap[rot].ch[op ^ 1];
	treap[rot].ch[op ^ 1] = treap[t].ch[op];
	treap[t].ch[op] = rot;
	rot = t;
	up(treap[rot].ch[op]);
	up(rot);
}
void insert(int& rot, int val)
{
	if (rot == 0)
	{
		rot = newnode(val);
		return;
	}
	if (treap[rot].val > val)
	{
		insert(treap[rot].ch[0], val);
		if (treap[rot].data < treap[treap[rot].ch[0]].data)
			rotate(rot, 1);
	}
	else if (treap[rot].val < val)
	{
		insert(treap[rot].ch[1], val);
		if (treap[rot].data < treap[treap[rot].ch[1]].data)
			rotate(rot, 0);
	}
	else
		treap[rot].cnt++;
	up(rot);
}
int find(int rot, int val)
{
	if (rot == 0)
		return -1;
	if (treap[rot].val == val)
		return rot;
	else if (treap[rot].val < val)
		return find(treap[rot].ch[1], val);
	else
		return find(treap[rot].ch[0], val);
}
int getpre(int rot, int val)//求比val小的最大的数字
{
	if (rot == 0)
		return 1;//最小值
	int ans = 1;//答案在bst中的下标
	if (treap[rot].val < val)
		ans = rot;
	if (treap[rot].val < val)
	{
		int t = getpre(treap[rot].ch[1], val);
		if (treap[ans].val < treap[t].val)
			ans = t;
	}
	else
	{
		int t = getpre(treap[rot].ch[0], val);
		if (treap[ans].val < treap[t].val)
			ans = t;
	}
	return ans;
}
int getnex(int rot, int val)//求比val大的最小的数字
{
	if (rot == 0)
		return 2;//最大值
	int ans = 2;
	if (treap[rot].val > val)
		ans = rot;
	if (treap[rot].val <= val)
	{
		int t = getnex(treap[rot].ch[1], val);
		if (treap[ans].val > treap[t].val)
			ans = t;
	}
	else
	{
		int t = getnex(treap[rot].ch[0], val);
		if (treap[ans].val > treap[t].val)
			ans = t;
	}
	return ans;
}
void delet(int& rot, int val)
{
	if (rot == 0)
		return;
	if (treap[rot].val == val)
	{
		if (treap[rot].cnt != 1)
			treap[rot].cnt--;
		else if (treap[rot].ch[0] || treap[rot].ch[1])
		{
			//判断左旋还是右旋，将需要删除的节点旋转到叶子节点
			if (treap[rot].ch[1] == 0 || treap[treap[rot].ch[0]].data > treap[treap[rot].ch[1]].data)
			{
				rotate(rot, 1);
				delet(treap[rot].ch[1], val);
			}
			else
			{
				rotate(rot, 0);
				delet(treap[rot].ch[0], val);
			}
		}
		else
			rot = 0;
	}
	else if (treap[rot].val > val)
		delet(treap[rot].ch[0], val);
	else
		delet(treap[rot].ch[1], val);
	up(rot);
}
int numrank(int rot, int val)
{
	if (rot == 0)
		return 0;
	if (treap[rot].val == val)
		return treap[treap[rot].ch[0]].sz + 1;
	else if (treap[rot].val > val)
		return numrank(treap[rot].ch[0], val);
	else
		return treap[treap[rot].ch[0]].sz + treap[rot].cnt + numrank(treap[rot].ch[1], val);
}
int kthnum(int rot, int k)
{
	if (k <= treap[treap[rot].ch[0]].sz)
		return kthnum(treap[rot].ch[0], k);
	else if (k > treap[treap[rot].ch[0]].sz + treap[rot].cnt)
		return kthnum(treap[rot].ch[1], k - (treap[treap[rot].ch[0]].sz + treap[rot].cnt));
	else
		return rot;
}
int main()
{
	int n;
	sc("%d", &n);
	build();
	while (n--)
	{
		int op, a;
		sc("%d%d", &op, &a);
		switch (op)
		{
		case 1:insert(root, a); break;
		case 2:delet(root, a); break;
		case 3:pr("%d\n", numrank(root, a) - 1); break;
		case 4:pr("%d\n", treap[kthnum(root, a + 1)].val); break;
		case 5:pr("%d\n", treap[getpre(root, a)].val); break;
		case 6:pr("%d\n", treap[getnex(root, a)].val); break;
		}
	}
}