bzoj 3224/Tyvj 1728 普通平衡树

  1 #include<cstdio>
  2 #include<cstdlib>
  3 #include<iostream>
  4 #include<algorithm>
  5 const int MAX_N = 100100;
  6 struct Node{
  7     int v, s, c;
  8     Node *ch[2];
  9     inline void set(int _v = 0, int _s = 0, Node *p = NULL){
 10         v = _v, s = c = _s;
 11         ch[0] = ch[1] = p;
 12     }
 13     inline void push_up(){
 14         s = ch[0]->s + ch[1]->s + c;
 15     }
 16     inline int cmp(int x) const{
 17         return x == v ? -1 : x > v;
 18     }
 19 };
 20 struct SizeBalanceTree{
 21     Node stack[MAX_N];
 22     Node *root, *null, *tail;
 23     Node *store[MAX_N];
 24     int top;
 25     void init(){
 26         tail = &stack[0];
 27         null = tail++;
 28         null->set();
 29         root = null;
 30         top = 0;
 31     }
 32     inline Node *newNode(int v){
 33         Node *p = null;
 34         if (top) p = store[--top];
 35         else p = tail++;
 36         p->set(v, 1, null);
 37         return p;
 38     }
 39     inline void rotate(Node* &x, int d){
 40         Node *k = x->ch[!d];
 41         x->ch[!d] = k->ch[d];
 42         k->ch[d] = x;
 43         k->s = x->s;
 44         x->push_up();
 45         x = k;
 46     }
 47     inline void Maintain(Node* &x, int d){
 48         if (x->ch[d] == null) return;
 49         if (x->ch[d]->ch[d]->s > x->ch[!d]->s){
 50             rotate(x, !d);
 51         } else if (x->ch[d]->ch[!d]->s > x->ch[!d]->s){
 52             rotate(x->ch[d], d), rotate(x, !d);
 53         } else {
 54             return;
 55         }
 56         Maintain(x, 0), Maintain(x, 1);
 57     }
 58     void insert(Node* &x, int v){
 59         if (x == null){
 60             x = newNode(v);
 61             return;
 62         } else {
 63             x->s++;
 64             int d = x->cmp(v);
 65             if (-1 == d){
 66                 x->c++;
 67                 return;
 68             }
 69             insert(x->ch[d], v);
 70             x->push_up();
 71             Maintain(x, d);
 72         }
 73     }
 74     void del(Node*  &x, int v){
 75         if (x == null) return;
 76         x->s--;
 77         int d = x->cmp(v);
 78         if (-1 == d){
 79             if (x->c > 1){
 80                 x->c--;
 81                 return;
 82             } else if (x->ch[0] == null || x->ch[1] == null){
 83                 store[top++] = x;
 84                 x = x->ch[0] == null ? x->ch[1] : x->ch[0];
 85             } else {
 86                 Node *ret = x->ch[1];
 87                 for (; ret->ch[0] != null; ret = ret->ch[0]);
 88                 del(x->ch[1], x->v = ret->v);
 89             }
 90         } else {
 91             del(x->ch[d], v);
 92         }
 93         if (x != null) x->push_up();
 94     }
 95     int find_kth(Node *x, int k){
 96         int t;
 97         for (; x != null;){
 98             t = x->ch[0]->s;
 99             if (k <= t) x = x->ch[0];
100             else if (t + 1 <= k && k <= t + x->c) break;
101             else k -= t + x->c, x = x->ch[1];
102         }
103         return x->v;
104     }
105     int find_rank(Node *x, int v){
106         int t, cur;
107         for (cur = 0; x != null;){
108             t = x->ch[0]->s;
109             if (v == x->v) break;
110             else if (v < x->v) x = x->ch[0];
111             else cur += t + x->c, x = x->ch[1];
112         }
113         return cur + t + 1;
114     }
115     int succ(Node *x, int v){
116         int ret = 0;
117         while (x != null){
118             if (x->v > v){
119                 ret = x->v;
120                 x = x->ch[0];
121             }
122             else x = x->ch[1];
123         }
124         return ret;
125     }
126     int pred(Node *x, int v){
127         int ret = 0;
128         while (x != null){
129             if (x->v < v){
130                 ret = x->v;
131                 x = x->ch[1];
132             }
133             else x = x->ch[0];
134         }
135         return ret;
136     }
137 }SBT;
138 int main(){
139 #ifdef LOCAL
140     freopen("in.txt", "r", stdin);
141     freopen("out.txt", "w+", stdout);
142 #endif
143     SBT.init();
144     int n, op, v;
145     scanf("%d", &n);
146     while (n--){
147         scanf("%d %d", &op, &v);
148         if (1 == op) SBT.insert(SBT.root, v);
149         else if (2 == op) SBT.del(SBT.root, v);
150         else if (3 == op) printf("%d\n", SBT.find_rank(SBT.root, v));
151         else if (4 == op) printf("%d\n", SBT.find_kth(SBT.root, v));
152         else if (5 == op) printf("%d\n", SBT.pred(SBT.root, v));
153         else  printf("%d\n", SBT.succ(SBT.root, v));
154      }
155     return 0;
156 }