标题说平衡树,那么应该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 push_up(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; push_up(r); push_up(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); push_up(r); } void erase(int& r, int v) { if (r == 0) return; if (v < node[r].val) erase(node[r].son[0], v); else if (v > node[r].val) erase(node[r].son[1], v); else { if (node[r].son[0] == 0 && node[r].son[1] == 0) { node[r].cnt --; node[r].size --; if (node[r].cnt == 0) r = 0; } else if (node[r].son[0] != 0 && node[r].son[1] == 0) { rotate(r, 1); erase(node[r].son[1], v); } else if (node[r].son[0] == 0 && node[r].son[1] != 0) { rotate(r, 0); erase(node[r].son[0], v); } 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); } } push_up(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\n", get_rank(root, x)); if (opt == 4) printf("%d\n", get_val(root, x)); if (opt == 5) printf("%d\n", get_pre(root, x)); if (opt == 6) printf("%d\n", get_suc(root, x)); } return 0; }
学到的第二个用随机数的算法,所以这个Treap的高度也是有的随机的。但是起码不太可能退化成链状。