[Treap]poj2985 The k-th Largest Group

题意

给出N只猫，然后有两种操作，分别是把两只猫所在的组合并，另一个是查询第k大的小组的猫的数量

思路

其实这道题目有很多种做法，但是查询第k大当然是用平衡树是最好的。那么在这里用并查集来记录每一个猫的小组，然后在Treap中插入每一个小组的数量。修改的时候就把原来的两个删去，然后再插入合并之后的值。
但是这种方法会超时……可能是Treap写得不好的原因。那么很容易想到优化，因为输入保证查询的是在范围内的，那么就可以利用一下溢出的情况。首先无视掉1的情况，只有合并时才会加入新的。那么如果查询查不到那么就证明答案是1，这个其实也很好想，只要画画样例就可以想到。

代码

之前写的时候一直TLE，然后就去写了个手写输入输出，结果RE了，最后重新把Treap打一遍才AC

#include 
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <algorithm>

using namespace std;

struct Treap {
    struct Node {
        Node *ch[2];
        int r;
        int v;
        int s;
        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;
            }
        }
    } *root;

    void init() {
        root = NULL;
//      srand(time(0));
    }

    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;
    }

    void insert(Node* &o, int x) {
        if(o == NULL) {
            o = new Node();
            o -> ch[0] = o -> ch[1] = NULL;
            o -> v = x;
            o -> r = rand();
        } else {
            int d = (x  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) {
            Node* u = o;
            if(o -> ch[0] != NULL && o -> ch[1] != NULL) {
                int d2 = (o -> ch[0] -> r > o -> ch[1] -> r ? 1 : 0);
                rotate(o, d2);
                remove(o -> ch[d2], x);
            } else {
                if(o -> ch[0] == NULL) {
                    o = o -> ch[1];
                } else {
                    o = o -> ch[0];
                }
                delete u;
            }
        } else {
            remove(o -> ch[d], x);
        }
        if(o != NULL) {
            o -> maintain();
        }
    }

    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;
    }

    int kth(Node* o, int k) {
        if(o == NULL || k > o -> s || k <= 0) {
            return 1;
        }
        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);
            }
        }
    }
} treap;

const int MAXN = 500000;

int siz[MAXN];
int fa[MAXN];

int _find(int root) {
    return fa[root] == root ? root : fa[root] = _find(fa[root]);
}

void Union(int a, int b) {
    int faA = _find(a);
    int faB = _find(b);
    if(faA != faB) {
        if(siz[faA] != 1) treap.remove(treap.root, siz[faA]);
        if(siz[faB] != 1) treap.remove(treap.root, siz[faB]);
        treap.insert(treap.root, siz[faA] + siz[faB]);
        fa[faA] = faB;
        siz[faB] += siz[faA];
    }
}

int NextInt() {
    char c = getchar();
    int x = 0;
    while(c < '0' || '9' < c) {
        c = getchar();
    }
    while('0' <= c && c <= '9') {
        x *= 10;
        x += c - '0';
        c = getchar();
    }
    return x;
}

int main(void) {
    int n, m;
    treap.init();
    //scanf("%d%d", &n, &m);
    n = NextInt();
    m = NextInt();
    for(int i = 1; i <= n; ++i) {
//      treap.insert(treap.root, 1);
        siz[i] = 1;
        fa[i] = i;
    }
    for(int i = 0; i < m; ++i) {
        int o;
        //scanf("%d", &o);
        o = NextInt();
        if(o == 1) {
            int k;
            //scanf("%d", &k);
            k = NextInt();
            printf("%d\n", treap.kth(treap.root, k));
        } else {
            int a, b;
            //scanf("%d%d", &a, &b);
            a = NextInt();
            b = NextInt();
            Union(a, b);
        }
    }
    return 0;
}