$Size Balanced Tree$

  1 #include <bits/stdc++.h>
  2 
  3 const int MaxN = 100010;
  4 
  5 int key[MaxN], lc[MaxN], rc[MaxN], size[MaxN];
  6 int root, cnt;
  7 
  8 void left_turn(int &x)
  9 {
 10     int y = rc[x];
 11     rc[x] = lc[y];
 12     lc[y] = x;
 13     size[y] = size[x];
 14     size[x] = size[lc[x]] + size[rc[x]] + 1;
 15     x = y;
 16 }
 17 
 18 void right_turn(int &x)
 19 {
 20     int y = lc[x];
 21     lc[x] = rc[y];
 22     rc[y] = x;
 23     size[y] = size[x];
 24     size[x] = size[lc[x]] + size[rc[x]] + 1;
 25     x = y;
 26 }
 27 
 28 void maintain(int &x, bool flag)
 29 {
 30     if (!flag)
 31     {
 32         if (size[lc[lc[x]]] > size[rc[x]])
 33             right_turn(x);
 34         else if (size[rc[lc[x]]] > size[rc[x]])
 35             left_turn(lc[x]), right_turn(x);
 36         else
 37             return;
 38     }
 39     else
 40     {
 41         if (size[rc[rc[x]]] > size[lc[x]])
 42             left_turn(x);
 43         else if (size[lc[rc[x]]] > size[lc[x]])
 44             right_turn(rc[x]), left_turn(x);
 45         else
 46             return;
 47     }
 48     maintain(lc[x], false);
 49     maintain(rc[x], true);
 50     maintain(x, true);
 51     maintain(x, false);
 52 }
 53 
 54 void insert(int &x, int val)
 55 {
 56     if (!x)
 57     {
 58         x = ++cnt;
 59         lc[x] = rc[x] = 0;
 60         size[x] = 1;
 61         key[x] = val;
 62     }
 63     else
 64     {
 65         size[x]++;
 66         if (val < key[x])
 67             insert(lc[x], val);
 68         else
 69             insert(rc[x], val);
 70         maintain(x, val >= key[x]);
 71     }
 72 }
 73 
 74 int del(int &x, int val)
 75 {
 76     size[x]--;
 77     if (val == key[x] || (val < key[x] && lc[x] == 0) || (val > key[x] && rc[x] == 0))
 78     {
 79         int y = key[x];
 80         if (lc[x] == 0 || rc[x] == 0)
 81             x = lc[x] + rc[x];
 82         else
 83             key[x] = del(lc[x], key[x] + 1);
 84         return y;
 85     }
 86     else
 87     {
 88         if (val < key[x])
 89             return del(lc[x], val);
 90         else
 91             return del(rc[x], val);
 92     }
 93 }
 94 
 95 int getmin()
 96 {
 97     int i;
 98     for (i = root; lc[i]; i = lc[i])
 99         ;
100     return key[i];
101 }
102 
103 int getmax()
104 {
105     int i;
106     for (i = root; rc[i]; i = rc[i])
107         ;
108     return key[i];
109 }
110 
111 int query_num(int &x, int rank)
112 {
113     if (size[lc[x]] + 1 == rank)
114         return key[x];
115     else if (rank <= size[lc[x]])
116         return query_num(lc[x], rank);
117     else
118         return query_num(rc[x], rank - size[lc[x]] - 1);
119 }
120 
121 int query_rank(int &x, int val)
122 {
123     if (!x)
124         return 1;
125     int ans = 0;
126     if (val <= key[x])
127         ans = query_rank(lc[x], val);
128     else
129         ans = size[lc[x]] + 1 + query_rank(rc[x], val);
130     return ans;
131 }
132 
133 int query_pre(int &x, int val)
134 {
135     if (!x)
136         return val;
137     int ans;
138     if (key[x] >= val)
139         ans = query_pre(lc[x], val);
140     else
141     {
142         ans = query_pre(rc[x], val);
143         if (ans == val)
144             ans = key[x];
145     }
146     return ans;
147 }
148 
149 int query_sub(int &x, int val)
150 {
151     if (!x)
152         return val;
153     int ans;
154     if (val >= key[x])
155         ans = query_sub(rc[x], val);
156     else
157     {
158         ans = query_sub(lc[x], val);
159         if (ans == val)
160             ans = key[x];
161     }
162     return ans;
163 }
164 
165 int main()
166 {
167     int n;
168     scanf("%d", &n);
169     root = cnt = 0;
170     for (int i = 1; i <= n; i++)
171     {
172         int op, x;
173         scanf("%d%d", &op, &x);
174         if (op == 1)
175             insert(root, x);
176         if (op == 2)
177             del(root, x);
178         if (op == 3)
179             printf("%d\n", query_rank(root, x));
180         if (op == 4)
181             printf("%d\n", query_num(root, x));
182         if (op == 5)
183             printf("%d\n", query_pre(root, x));
184         if (op == 6)
185             printf("%d\n", query_sub(root, x));
186     }
187     return 0;
188 }