Size Balanced Tree

  1 #include <cstdio>
  2 
  3 const int maxn = 1000000;
  4 
  5 int n, a, b, c, root;
  6 int Child[maxn][2];
  7 int Key[maxn], Size[maxn], Count;
  8 
  9 int New_node (int v)
 10 {
 11     Count++;
 12     Key[Count] = v;
 13     Size[Count] = 1;
 14     return Count;
 15 }
 16 
 17 void Rotate (int& x, int d)
 18 {
 19     int k = Child[x][!d];
 20     Child[x][!d] = Child[k][d];
 21     Child[k][d] = x;
 22     Size[k] = Size[x];
 23     Size[x] = Size[Child[x][0]] + Size[Child[x][1]] + 1;
 24     x = k;
 25 }
 26 
 27 void Maintain (int& t, int d)
 28 {
 29     if (!t) return;
 30     if (Size[Child[Child[t][d]][d]] > Size[Child[t][!d]])
 31         Rotate (t, !d);
 32     else if (Size[Child[Child[t][d]][!d]] > Size[Child[t][!d]])
 33         Rotate (Child[t][d], d), Rotate (t, !d);
 34     else return;
 35     Maintain (Child[t][0], 0);
 36     Maintain (Child[t][1], 1);
 37     Maintain (t, 0);
 38     Maintain (t, 1);
 39 }
 40 
 41 void Insert (int& t, int v)
 42 {
 43     if (t == 0) t = New_node (v);
 44     else
 45     {
 46         Size[t]++;
 47         Insert (Child[t][v >= Key[t]], v);
 48         Maintain (t, v >= Key[t]);
 49     }
 50 }
 51 
 52 int Delete (int& t, int v)
 53 {
 54     int Ret;
 55     Size[t]--;
 56     if (v == Key[t] || Child[t][v > Key[t]] == 0)
 57     {
 58         Ret = Key[t];
 59         if (Child[t][0] && Child[t][1])
 60             Key[t] = Delete (Child[t][0], v + 1);
 61         else t = Child[t][0] + Child[t][1];
 62     }else
 63         Ret = Delete (Child[t][v > Key[t]], v);
 64     return Ret;
 65 }
 66 
 67 int Rank (int t, int v)
 68 {
 69     if (t == 0) return 1;
 70     if (v <= Key[t]) return Rank (Child[t][0], v);
 71     return Rank (Child[t][1], v) + Size[Child[t][0]] + 1;
 72 }
 73 
 74 int Select (int t, int k)
 75 {
 76     if (Child[t][0] && Size[Child[t][0]] >= k)
 77         return Select (Child[t][0], k);
 78     if (Size[Child[t][0]] + 1 == k) return Key[t];
 79     return Select (Child[t][1], k - Size[Child[t][0]] - 1);
 80 }
 81 
 82 void Print (int t)
 83 {
 84     if (t == 0) return;
 85     Print (Child[t][0]);
 86     printf ("%d ", Key[t]);
 87     Print (Child[t][1]);
 88 }
 89 
 90 int main ()
 91 {
 92     scanf ("%d", &n);
 93     for (int i = 0; i < n; i++)
 94     {
 95         scanf ("%d %d", &a, &b);
 96         if (a == 1) Insert (root, b);
 97         if (a == 2)
 98         {
 99             c = Delete (root, b);
100             if (c != b) Insert (root, c);
101         }
102         if (a == 3) printf ("%d\n", Select (root, b));
103         if (a == 4) printf ("%d\n", Rank (root, b));
104         if (a == 5)
105         {
106             if (b == 1)
107             {
108                 for (int j = 1; j <= Count; j++)
109                     printf ("Child : %d %d, Key : %d\n", Child[j][0], Child[j][1], Key[j]);
110             }else Print (root), printf ("\n");
111         }
112     }
113 }