题目传送门
标题说平衡树，那么应该AVL，红黑树都能过，但是这次做这题主要是学习Treap，所以花了几天搞出了这题。其他方法以后再说吧
- Treap（带旋转）
#include <bits/stdc++.h> using namespace std; const int MAXN = 100010; const int INF = 0x3f3f3f3f; struct Treap { int son[2]; int val, rand; int cnt, size; } node[MAXN]; int root, tot; void update(int i) { int ls = node[i].son[0]; int rs = node[i].son[1]; node[i].size = node[ls].size + node[rs].size + node[i].cnt; } void rotate(int& r, int index) { int k = node[r].son[index ^ 1]; node[r].son[index ^ 1] = node[k].son[index]; node[k].son[index] = r; update(r); update(k); r = k; } void insert(int& r, int v) { if (r == 0) { r = ++ tot; node[r].val = v; node[r].rand = rand(); node[r].size = node[r].cnt = 1; return; } if (node[r].val == v) { node[r].size ++; node[r].cnt ++; return; } int index = v > node[r].val; insert(node[r].son[index], v); if (node[node[r].son[index]].rand < node[r].rand) rotate(r, index ^ 1); update(r); } void erase(int& r, int v) { if (r == 0) return; if (v == node[r].val) { if (node[r].cnt > 1) { node[r].cnt --; node[r].size --; return; } if (node[r].son[0] == 0 || node[r].son[1] == 0) { r = node[r].son[0] + node[r].son[1]; } else { int ls = node[r].son[0]; int rs = node[r].son[1]; int index = node[ls].rand < node[rs].rand; rotate(r, index); erase(node[r].son[index], v); } } else { int index = v > node[r].val; erase(node[r].son[index], v); } update(r); } int get_rank(int r, int v) { int size = node[node[r].son[0]].size; int cnt = node[r].cnt; if (node[r].val == v) return size + 1; if (node[r].val > v) return get_rank(node[r].son[0], v); if (node[r].val < v) return size + cnt + get_rank(node[r].son[1], v); } int get_val(int r, int rk) { int size = node[node[r].son[0]].size; int cnt = node[r].cnt; if (rk < size + 1) return get_val(node[r].son[0], rk); if (rk >= size + 1 && rk <= size + cnt) return node[r].val; if (rk > size + cnt) return get_val(node[r].son[1], rk - size - cnt); } int get_pre(int r, int v) { if (r == 0) return -INF; if (node[r].val >= v) return get_pre(node[r].son[0], v); else return max(node[r].val, get_pre(node[r].son[1], v)); } int get_suc(int r, int v) { if (r == 0) return INF; if (node[r].val <= v) return get_suc(node[r].son[1], v); else return min(node[r].val, get_suc(node[r].son[0], v)); } int main() { srand(time(NULL)); int n; scanf("%d", &n); for (int i = 1; i <= n; i++) { int opt, x; scanf("%d%d", &opt, &x); if (opt == 1) insert(root, x); if (opt == 2) erase(root, x); if (opt == 3) printf("%d ", get_rank(root, x)); if (opt == 4) printf("%d ", get_val(root, x)); if (opt == 5) printf("%d ", get_pre(root, x)); if (opt == 6) printf("%d ", get_suc(root, x)); } return 0; }
学到的第二个用随机数的算法，所以这个Treap的高度也是有的随机的。但是起码不太可能退化成链状。
- Treap（无旋转）
#include <bits/stdc++.h> using namespace std; const int MAXN = 100010; const int INF = 0x3f3f3f3f; struct Treap { int ls, rs; int val, rand; int size; } node[MAXN]; int tot, root; int add_node(int v) { int i = ++tot; node[i].val = v; node[i].rand = rand(); node[i].ls = node[i].rs = 0; node[i].size = 1; return i; } void split(int rt, int& a, int& b, int v) { if (rt == 0) { a = b = 0; return; } if (node[rt].val <= v) { a = rt; split(node[rt].rs, node[a].rs, b, v); } else { b = rt; split(node[rt].ls, a, node[b].ls, v); } node[rt].size = node[node[rt].ls].size + node[node[rt].rs].size + 1; } void merge(int& rt, int a, int b) { if (a == 0 || b == 0) { rt = a + b; return; } if (node[a].rand < node[b].rand) { rt = a; merge(node[rt].rs, node[a].rs, b); } else { rt = b; merge(node[rt].ls, a, node[b].ls); } node[rt].size = node[node[rt].ls].size + node[node[rt].rs].size + 1; } void insert(int& rt, int v) { int x = 0, y = 0; split(rt, x, y, v); merge(x, x, add_node(v)); merge(rt, x, y); } void erase(int&rt, int v) { int x = 0, y = 0, z = 0; split(rt, x, z, v); split(x, x, y, v - 1); merge(y, node[y].ls, node[y].rs); merge(x, x, y); merge(rt, x, z); } int get_rank(int rt, int v) { int x = 0, y = 0; split(rt, x, y, v - 1); int rk = node[x].size + 1; merge(rt, x, y); return rk; } int get_val(int rt, int rk) { int size = node[node[rt].ls].size; if (rk < size + 1) return get_val(node[rt].ls, rk); if (rk == size + 1) return node[rt].val; if (rk > size + 1) return get_val(node[rt].rs, rk - size - 1); } int get_pre(int rt, int v) { int x = 0, y = 0; split(rt, x, y, v - 1); int pre = get_val(x, node[x].size); merge(rt, x, y); return pre; } int get_suc(int rt, int v) { int x = 0, y = 0; split(rt, x, y, v); int suc = get_val(y, 1); merge(rt, x, y); return suc; } int main() { srand(time(NULL)); int n; scanf("%d", &n); for (int i = 1; i <= n; i++) { int opt, x; scanf("%d%d", &opt, &x); if (opt == 1) insert(root, x); if (opt == 2) erase(root, x); if (opt == 3) printf("%d ", get_rank(root, x)); if (opt == 4) printf("%d ", get_val(root, x)); if (opt == 5) printf("%d ", get_pre(root, x)); if (opt == 6) printf("%d ", get_suc(root, x)); } return 0; }
昨天看懂了带旋转的Treap，今天来写无旋转的Treap，感觉确实是比带旋转的好敲而且好懂，关键就是理解split()和merge()两个核心函数。