POJ2985 The k-th Largest Group treap

POJ2985 比较简单的平衡树题目 树内不要添加容量为1的节点 否则会超时。

#include<iostream>
#include<cstdio>
#include<cstdlib>
#include<cmath>
#include<cstring>
#include<vector>
#include<queue>
#include<algorithm>
using namespace std;
const int maxn=500000+10,INF=1e+7;
int root,treapcnt,group[maxn],num[maxn],key[maxn],priority[maxn],child[maxn][2],size[maxn];
int treaptot=0;
inline int cmp(int a,int b)
{
 if(num[a]!=num[b])return num[a]<num[b];
 if(a!=b)return a<b;
 return -1;
}
void Treap()
{
 root=0;
 treapcnt=1;
 priority[0]=INF;
 size[0]=0;
}

inline void update(int x)
{
 size[x]=size[child[x][0]]+1+size[child[x][1]];
}

void rotate(int &x,int t)
{
 int y=child[x][t];
 child[x][t]=child[y][1-t];
 child[y][1-t]=x;
 update(x);update(y);
 x=y;
}

void _insert(int &x,int k)
{
 if(x)
 	{
 	 if(cmp(key[x],k)==-1)
 	 	{
 	 	 return ;
		}
		else
			{
			 int t=cmp(key[x],k);
			 _insert(child[x][t],k);
			 if(priority[child[x][t]]<priority[x])rotate(x,t);
			}
	}
	else
		{
		 x=treapcnt++;
		 key[x]=k;
		 priority[x]=rand();
		 child[x][0]=child[x][1]=0;
		}
 update(x);
}

void _erase(int &x,int k)
{
 int flag=cmp(key[x],k);
 if(flag==-1)
 	{

	     		 if(child[x][0]==0&&child[x][1]==0){x=0;return ;}
	     		 int t=priority[child[x][0]]>priority[child[x][1]];
	     		 rotate(x,t);
	     		 _erase(x,k);

	}else
		{
		 _erase(child[x][flag],k);
		}
 update(x);
}

int _getKth(int &x,int k)
{
 if(k>size[x])return -1;
 if(k<=size[child[x][0]])return _getKth(child[x][0],k);
 k-=size[child[x][0]]+1;
 if(k<=0)return key[x];
 return _getKth(child[x][1],k);
}
inline int find(int x)
{
 if(0==group[x])return x;
 	else return group[x]=find(group[x]);
}
inline void combine(int a,int b)
{
 a=find(a);b=find(b);
 if(a==b)return ;
 if(num[a]>1){treaptot--;_erase(root,a);}
 if(num[b]>1){treaptot--;_erase(root,b);}
 group[b]=a;
 num[a]+=num[b];
 treaptot++;_insert(root,a);
}
int main()
{freopen("t.txt","r",stdin);
 freopen("1.txt","w",stdout);
 srand(100717);
 int n,m;
 Treap();
 scanf("%d%d",&n,&m);
 for(int i=1;i<=n;i++)
 	{num[i]=1;group[i]=0;}
 int flag,a,b;
 for(int i=1;i<=m;i++)
 	{

 	 scanf("%d",&flag);
 	 if(flag)
 	 	{
 	 	 scanf("%d",&a);
 	 	 if(a>size[root])printf("1\n");
 	 		else  printf("%d\n",num[_getKth(root,treaptot-a+1)]);
		}
		else
			{
			 n--;
			 scanf("%d%d",&a,&b);
			 combine(a,b);
			}
	}

 return 0;
}