Treap实现名次树

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 const int maxn = 1000 + 10;
  4 
  5 struct node {
  6     node *ch[2];
  7     int r;
  8     int v;
  9     int s;
 10     node (int v):v(v) {ch[0] = ch[1] = NULL; r = rand(); s = 1;}
 11     bool operator < (const node& rhs) const {
 12         return r < rhs.r;
 13     }
 14     int cmp(int x) const{
 15         if (x == v) return -1;
 16         return x < v ? 0 : 1;
 17     }
 18     void maintain() {
 19         s = 1;
 20         if (ch[0] != NULL) s += ch[0]->s;
 21         if (ch[1] != NULL) s += ch[1]->s;
 22     }
 23 };
 24 
 25 void rotate(node* &o, int d) {
 26     node* k = o->ch[d^1]; o->ch[d^1] = k->ch[d]; 
 27     k->ch[d] = o; o->maintain(); k->maintain(); o = k;
 28 }
 29 
 30 int find(node* o, int x) {
 31     while(o != NULL) {
 32         int d = o->cmp(x);
 33         if (d == -1) return 1;
 34         else o = o->ch[d];
 35     }
 36     return 0;
 37 }
 38 
 39 void insert(node* &o, int x) {
 40     if (o == NULL) {
 41         o = new node(x);
 42         o->ch[0] = o->ch[1] = NULL;
 43         o->v = x;
 44     }
 45     else {
 46         int d = (x < o->v ? 0 : 1);
 47         insert(o->ch[d], x);
 48         if (o->ch[d]->r > o->r) rotate(o, d^1);
 49     }
 50     o->maintain();
 51 }
 52 
 53 void remove(node* &o, int x) {
 54     int d = o->cmp(x);
 55     if (d == -1) {
 56         if (o->ch[0] == NULL) o = o->ch[1];
 57         else if (o->ch[1] == NULL) o = o->ch[0];
 58         else {
 59             int d2 = o->ch[0]->r > o->ch[1]->r ? 1 : 0;
 60             rotate(o, d2); remove(o->ch[d2], x);
 61         }
 62     }
 63     else remove(o->ch[d], x); 
 64     if (o != NULL) o->maintain();
 65 }
 66 
 67 int kth(node* &o, int k) {
 68     if (o == NULL || k <= 0 || k > o->s) return 0;
 69     int s = (o->ch[1] == NULL ? 0 : o->ch[1]->s);
 70     if (k == s + 1) return o->v;
 71     else if (k <= s) return kth(o->ch[1], k);
 72     else return kth(o->ch[0], k-s-1);
 73 }
 74 
 75 int main() {
 76     int a, k;
 77     char op[10];
 78     node *rt = NULL;
 79     while(cin >> op, op[0] != ‘E‘) {
 80         switch(op[0]) {
 81             case ‘i‘:
 82                 cin >> a;
 83                 if (find(rt, a)) {
 84                     cout << a << " is in Treap" << endl;
 85                     break;
 86                 }
 87                 insert(rt, a);
 88                 break;
 89             case ‘r‘:
 90                 cin >> a;
 91                 if (!find(rt, a)) {
 92                     cout << a << " isn‘t in Treap" << endl;
 93                     break;
 94                 }
 95                 remove(rt, a);
 96                 break;
 97             case ‘q‘:
 98                 cin >> k;
 99                 a = kth(rt, k);
100                 if (a) cout << "The " << k << "th is " << a << endl;
101                 else cout << "error" << endl;
102                 break;
103         }
104     }
105 
106     return 0;
107 }