P3369 普通平衡树模板 treap

/problem/P3369

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

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

///problem/P3369//treap#include <bits/stdc++.h>#define ll long long#define sc scanf#define pr printfusing 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);}elsetreap[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);elsereturn 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);}}elserot = 0;}else if (treap[rot].val > val)delet(treap[rot].ch[0], val);elsedelet(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);elsereturn 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));elsereturn 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;}}}