POJ2985 The k-th Largest Group(treap+并查集)

题目链接

Input

1st line: Two numbers N and M (1 ≤ N, M ≤ 200,000), namely the number of cats and the number of operations.

2nd to (m + 1)-th line: In each line, there is number C specifying the kind of operation Newman wants to do. If C = 0, then there are two numbers i and j (1 ≤ i, j ≤ n) following indicating Newman wants to combine the group containing the two cats (in case these two cats are in the same group, just do nothing); If C = 1, then there is only one number k (1 ≤ k ≤ the current number of groups) following indicating Newman wants to know the size of the k-th largest group.

Output

For every operation “1” in the input, output one number per line, specifying the size of the kth largest group.

Sample Input

10 10
0 1 2
1 4
0 3 4
1 2
0 5 6
1 1
0 7 8
1 1
0 9 10
1 1

Sample Output

1
2
2
2
2

题目大意：给定N只猫，序号从一到N，一开始每个都是一组，然后对这群猫进行操作，可以把两只猫合成一组，然后可以查询第K大的组里有几只猫。

并查集加treap,对于每一次合并，在treap里分别将两元素删除，然后用并查集将其合并，在将合并后的结果加入treap树即可。

代码如下：

#include <iostream>
#include <string.h>
#include <stdlib.h>
#include <algorithm>
#include <stdio.h>
using namespace std;
const int N=200005;
struct node
{
    int l;
    int r;
    int val;
    int ran;
    int num;
    int w;
}tree[N];
struct query
{
    int c;
    int a,b;
}in[N];

int pre[N];//并查集
int s,root;

int fin(int x)//并查集查找
{
    if(x!=pre[x])
        pre[x]=fin(pre[x]);
    return pre[x];
}

void join(int x,int y)//并查集合并函数
{
    int fx=fin(x),fy=fin(y);
    if(fx!=fy)
    {
        pre[fy]=fx;
    }
}

void lturn(int &k)
{
    int t=tree[k].r;
    tree[k].r=tree[t].l;
    tree[t].l=k;
    tree[t].num=tree[k].num;
    tree[k].num=tree[tree[k].l].num+tree[tree[k].r].num+tree[k].w;
    k=t;
}

void rturn(int &k)
{
    int t=tree[k].l;
    tree[k].l=tree[t].r;
    tree[t].r=k;
    tree[t].num=tree[k].num;
    tree[k].num=tree[tree[k].l].num+tree[tree[k].r].num+tree[k].w;
    k=t;
}

void ins(int &k,int x)
{
    if(k==0)
    {
        s++;
        k=s;
        tree[k].val=x;
        tree[k].num=tree[k].w=1;
        tree[k].ran=rand();
        return ;
    }
    tree[k].num++;
    if(tree[k].val==x)
        tree[k].w++;
    else if(tree[k].val>x)
    {
        ins(tree[k].l,x);
        if(tree[tree[k].l].ran<tree[k].ran)
            rturn(k);
    }
    else
    {
        ins(tree[k].r,x);
        if(tree[tree[k].r].ran<tree[k].ran)
            lturn(k);
    }
}
void del(int &k,int x)
{
    if(k==0)
        return;
    if(tree[k].val==x)
    {
        if(tree[k].w>1)
        {
            tree[k].num--;
            tree[k].w--;
            return;
        }
        if(tree[k].l*tree[k].r==0)
            k=tree[k].l+tree[k].r;
        else if(tree[tree[k].l].ran<tree[tree[k].r].ran)
        {
            rturn(k);
            del(k,x);
        }
        else
        {
            lturn(k);
            del(k,x);
        }
    }
    else if(tree[k].val<x)
    {
        tree[k].num--;
        del(tree[k].r,x);
    }
    else
    {
        tree[k].num--;
        del(tree[k].l,x);
    }
}
int queryknum(int k,int x)
{
    if(k==0)
        return 0;
    if(x<=tree[tree[k].l].num)
        return queryknum(tree[k].l,x);
    else if(x>tree[tree[k].l].num+tree[k].w)
        return queryknum(tree[k].r,x-tree[tree[k].l].num-tree[k].w);
    else
        return tree[k].val;
}
int main()
{
    int data[N];
    int n,m;
    memset(tree,0,sizeof(tree));
    memset(data,0,sizeof(data));
    memset(in,0,sizeof(in));
    memset(pre,0,sizeof(pre));
    scanf("%d %d",&n,&m);
    root=0;
    s=0;
    for(int i=1;i<=n;i++)
    {
        pre[i]=i;
        data[i]=1;
        ins(root,1);
    }
    for(int i=1;i<=m;i++)
    {
        scanf("%d",&in[i].c);
        if(in[i].c==0)
            scanf("%d %d",&in[i].a,&in[i].b);
        else
            scanf("%d",&in[i].a);
    }
    int total=n;
    for(int i=1;i<=m;i++)
    {
        if(in[i].c==0)
        {
            int x=fin(in[i].a);
            int y=fin(in[i].b);
            if(x!=y)
            {
                total--;//两元素合并，树内的总元素数目减1
                del(root,data[x]);
                del(root,data[y]);
                join(x,y);//x是y的父节点
                data[x]+=data[y];//更改后的元素
                ins(root,data[x]);
            }
        }
        if(in[i].c==1)
        {
            printf("%d\n",queryknum(root,total-in[i].a+1));//查询为第K大的元素，所以倒过来查找，将treap插入过程的方向更改也可
        }
    }
    return 0;
}