Treap实现名次树

Treap（树堆）的大部分功能STL的set都可以实现，但因为set的过度封装使得某些特定的功能不能实现，比如求第k大的值。

Code:

#include<bits/stdc++.h>
using namespace std;
const int maxn = 1000 + 10;

struct node {
    node *ch[2];
    int r;
    int v;
    int s;
    node (int v):v(v) {ch[0] = ch[1] = NULL; r = rand(); s = 1;}
    bool operator < (const node& rhs) const {
        return r < rhs.r;
    }
    int cmp(int x) const{
        if (x == v) return -1;
        return x < v ? 0 : 1;
    }
    void maintain() {
        s = 1;
        if (ch[0] != NULL) s += ch[0]->s;
        if (ch[1] != NULL) s += ch[1]->s;
    }
};

void rotate(node* &o, int d) {
    node* k = o->ch[d^1]; o->ch[d^1] = k->ch[d]; 
    k->ch[d] = o; o->maintain(); k->maintain(); o = k;
}

int find(node* o, int x) {
    while(o != NULL) {
        int d = o->cmp(x);
        if (d == -1) return 1;
        else o = o->ch[d];
    }
    return 0;
}

void insert(node* &o, int x) {
    if (o == NULL) {
        o = new node(x);
        o->ch[0] = o->ch[1] = NULL;
        o->v = x;
    }
    else {
        int d = (x < o->v ? 0 : 1);
        insert(o->ch[d], x);
        if (o->ch[d]->r > o->r) rotate(o, d^1);
    }
    o->maintain();
}

void remove(node* &o, int x) {
    int d = o->cmp(x);
    if (d == -1) {
        if (o->ch[0] == NULL) o = o->ch[1];
        else if (o->ch[1] == NULL) o = o->ch[0];
        else {
            int d2 = o->ch[0]->r > o->ch[1]->r ? 1 : 0;
            rotate(o, d2); remove(o->ch[d2], x);
        }
    }
    else remove(o->ch[d], x); 
    if (o != NULL) o->maintain();
}

int kth(node* &o, int k) {
    if (o == NULL || k <= 0 || k > o->s) return 0;
    int s = (o->ch[1] == NULL ? 0 : o->ch[1]->s);
    if (k == s + 1) return o->v;
    else if (k <= s) return kth(o->ch[1], k);
    else return kth(o->ch[0], k-s-1);
}

int main() {
    int a, k;
    char op[10];
    node *rt = NULL;
    while(cin >> op, op[0] != 'E') {
        switch(op[0]) {
            case 'i':
                cin >> a;
                if (find(rt, a)) {
                    cout << a << " is in Treap" << endl;
                    break;
                }
                insert(rt, a);
                break;
            case 'r':
                cin >> a;
                if (!find(rt, a)) {
                    cout << a << " isn't in Treap" << endl;
                    break;
                }
                remove(rt, a);
                break;
            case 'q':
                cin >> k;
                a = kth(rt, k);
                if (a) cout << "The " << k << "th is " << a << endl;
                else cout << "error" << endl;
                break;
        }
    }

    return 0;
}