平衡树

Posted 2021-02-08 ac-evil

$Treap$实现

  1 #include <bits/stdc++.h>
  2 
  3 using namespace std;
  4 
  5 #define re register
  6 #define rep(i, a, b) for (re int i = a; i <= b; ++i)
  7 #define repd(i, a, b) for (re int i = a; i >= b; --i)
  8 #define maxx(a, b) a = max(a, b);
  9 #define minn(a, b) a = min(a, b);
 10 #define LL long long
 11 #define INF (1 << 30)
 12 
 13 inline int read() {
 14     int w = 0, f = 1; char c = getchar();
 15     while (!isdigit(c)) f = c == ‘-‘ ? -1 : f, c = getchar();
 16     while (isdigit(c)) w = (w << 3) + (w << 1) + (c ^ ‘0‘), c = getchar();
 17     return w * f;
 18 }
 19 
 20 struct Node {
 21     Node* ch[2];
 22     int v, r, c, s;
 23     int cmp(int x) const {
 24         if (x == v) return -1;
 25         return x < v ? 0 : 1;
 26     }
 27     void maintain() {
 28         s = c + ch[0]->s + ch[1]->s;
 29     }
 30 };
 31 
 32 Node* null;
 33 
 34 void newnode(Node* &o) { o = new Node(); o->r = rand(); o->ch[0] = o->ch[1] = null; o->v = o->s = o->c = 0; }
 35 void setnode(Node* &o, int v) { o->v = v; o->s = o->c = 1; }
 36 
 37 void rotate(Node* &o, int d) {
 38     Node* k = o->ch[d^1]; o->ch[d^1] = k->ch[d]; k->ch[d] = o;
 39     o->maintain(), k->maintain(); o = k;
 40 }
 41 
 42 void insert(Node* &o, int v) {
 43     if (o == null) { newnode(o); setnode(o, v); return; }
 44     int d = o->cmp(v);
 45     if (d == -1) o->c++, o->s++;
 46     else {
 47         insert(o->ch[d], v);
 48         if (o->ch[d]->r > o->r) rotate(o, d^1);
 49         o->maintain();
 50     }
 51 }
 52 
 53 void remove(Node* &o, int v) {
 54     int d = o->cmp(v);
 55     if (d == -1) {
 56         if (o->c) o->c--;
 57         if (!o->c) {
 58             if (o->ch[0] != null && o->ch[1] != null) {
 59                 int d2 = o->ch[0]->r < o->ch[1]->r ? 0 : 1;
 60                 rotate(o, d2); remove(o->ch[d2], v);
 61             }
 62             else {
 63                 Node* u = o;
 64                 if (o->ch[0] != null) o = o->ch[0]; else o = o->ch[1];
 65                 delete u;
 66             }
 67         }
 68     }
 69     else remove(o->ch[d], v);
 70     o->maintain();
 71 }
 72 
 73 int get_rank(Node* o, int v) {
 74     int d = o->cmp(v);
 75     if (d == -1) return o->ch[0]->s + 1;
 76     return get_rank(o->ch[d], v) + d * (o->ch[0]->s + o->c);
 77 }
 78 
 79 int get_val(Node* o, int rk) {
 80     if (rk <= o->ch[0]->s) return get_val(o->ch[0], rk);
 81     else if (rk > o->ch[0]->s + o->c) return get_val(o->ch[1], rk - o->ch[0]->s - o->c);    
 82     else return o->v;
 83 }
 84 
 85 int get_pre(Node* o, int v) {
 86     if (o == null) return -INF;
 87     if (o->v >= v) return get_pre(o->ch[0], v);
 88     else return max(o->v, get_pre(o->ch[1], v));
 89 }
 90 
 91 int get_next(Node* o, int v) {
 92     if (o == null) return INF;
 93     if (o->v <= v) return get_next(o->ch[1], v);
 94     else return min(o->v, get_next(o->ch[0], v));
 95 }
 96 
 97 Node* root;
 98 
 99 int n;
100 
101 int main() {
102     newnode(null);
103     root = null;
104     n = read();
105     rep(i, 1, n) {
106         int opt = read(), v = read();
107         if (opt == 1) insert(root, v);
108         if (opt == 2) remove(root, v);
109         if (opt == 3) printf("%d
", get_rank(root, v));
110         if (opt == 4) printf("%d
", get_val(root, v));
111         if (opt == 5) printf("%d
", get_pre(root, v));
112         if (opt == 6) printf("%d
", get_next(root, v));
113     }
114     return 0;
115 }