PKU 2985 The k-th Largest Group

http://poj.org/problem?id=2985
n匹の猫がいて、最初は1匹が1グループに対応している状態から、ある猫がいるグループとある猫がいるグループを併合するというクエリと
全体でk番目に大きいグループの大きさを尋ねうクエリを処理するという問題。

さきほどの問題のように、グループの大きさをユニオンファインドしながら管理しつつ、出現する大きさをセグメント木で増やしたり、減らしたりしてクエリに答える。

配列のサイズ確保ミスで間違ったり、同じグループの猫がきたら無視するというのを読み飛ばしたりして、ちょっとWAとかもらいました。

int un[300000];
int ch[300000];

int find(int x){
  if(x==un[x])
    return x;
  return
    un[x] = find(un[x]);
}

void unit(int a,int b){
  int ap=find(a);
  int bp=find(b);
  un[ap] = bp;
  ch[bp] += ch[ap];
}

int seg[1000000];

void update(int cl,int cr,int k,int p,int x){
  if(p<cl || cr<=p) return;
  if(cl==p && p+1==cr){
    seg[k] += x;
    return;
  }

  update(cl,cl+cr>>1,k*2+1,p,x);
  update(cl+cr>>1,cr,k*2+2,p,x);
  seg[k] = seg[k*2+1] + seg[k*2+2];
}

int query(int cl,int cr,int k,int q){
  if(cl+1==cr) return cl;
  if(seg[k*2+2]>=q)
    return query(cl+cr>>1,cr,k*2+2,q);
  return
    query(cl,cl+cr>>1,k*2+1,q-seg[k*2+2]);
}

main(){
  int n,m;
  scanf("%d%d",&n,&m);
  //cin >> n >> m;
  ++n;  
  rep(i,n){
    un[i] = i;
    ch[i] = 1;
  }

  update(0,n,0,1,n);

  rep(i,m){
    int q;
    scanf("%d",&q);
    if(q) {
      int k;
      scanf("%d",&k);
      
      printf("%d\n" ,query(0,n,0,k) );
    }else{
      int a,b;
      scanf("%d%d",&a,&b);
      --a,--b;
      int pa,pb;
      pa=find(a);
      pb=find(b);
      if(pa==pb) continue;
      update(0,n,0,ch[pa],-1);
      update(0,n,0,ch[pb],-1);
      update(0,n,0,ch[pa]+ch[pb],1);
      unit(a,b);
    }
  }
}