Official

F - BOX Editorial by physics0523


タイプ \(1\) の操作は集合同士を併合する操作、すなわち Union-Find で実行可能な操作そのものです。なので、 Union-Find を活用してこの問題を解くという方針が立ちます。

まず、ボールについては高々 \(N+Q\) 番のものまでしか考えなくてよいです。また、タイプ \(1\) の操作はボールの集合同士の併合なので、 \(N+Q\) 要素の Union-Find を用意しておけばボールの追加への対応が十分可能そうです。

問題となるのはタイプ \(3\) の操作です。「ボールが沢山集まっている集合」が入っている箱が目まぐるしく変わる可能性がある (例えば、 \(10^5\) 個程度のボールを箱 \(1\) に集めてそれらを箱 \(2,3,\dots\) へ動かしていくような操作が行われる場合) ので、「各ボールがどの箱に入っているか」の情報を毎回付け替えていては TLE してしまいます。
そこで、「どのボールの集合がどの箱に入っているか」という形で情報を保持したいという方針が立ちます。各ボールごとではなくボールの集合でまとめてボールの移動を行えば計算量を少なく抑えられそうです。
では、この方針を実現していきましょう。

鍵となるのは、 Union-Find の代表元を返す機能です。 ACL でいう dsu の leader の機能 です。

以下の情報を管理していくことを考えます。

  • \(ld[x] = \)\(x\) に入っている集合の代表元。但し、箱 \(x\) が空である場合は \(-1\) とする。
  • \(inv[y] = \) 代表元が \(y\) である集合がどの箱に入っているか。但し、代表元が \(y\) である集合がない場合の値は不定である。

タイプ \(1\) :

  • \(Y\) が空であれば、何もしなくてよい。
  • そうでなく箱 \(X\) が空であれば、以下の手順で箱 \(Y\) の中身を箱 \(X\) に移すことができる。
    • \(ld[X]=ld[Y]\) とする。これは箱 \(Y\) に入っていた集合を箱 \(X\) に移動させることに対応する。
    • \(inv[ld[X]]=X\) とする。
    • \(ld[Y]=-1\) とする。これは箱 \(Y\) を空にすることに対応する。
  • そうでない場合、 \(X,Y\) 双方に少なくとも \(1\) 個のボールが入っているので、以下の手順を踏む。
    • 代表元が \(ld[X]\) であるボールの集合と代表元が \(ld[Y]\) であるボールの集合とを併合し、その集合の代表元を新たに \(ld[X]\) とする。
    • \(inv[ld[X]]=X\) とする。
    • \(ld[Y]=-1\) とする。

タイプ \(2\) :
新たに入れるボールの番号を \(b\) とする。

  • \(X\) が空なら、 \(ld[X]=b, inv[ld[X]]=X\) とする。
  • そうでない場合、代表元が \(ld[X]\) であるボールの集合と代表元が \(b\) であるボールの集合とを併合し、その集合の代表元を新たに \(ld[X]\) とする。その後、 \(inv[ld[X]]=X\) とする。

タイプ \(3\) :
ボール \(X\) が含まれるボールの集合の代表元を \(tg\) とすると、 \(inv[tg]\) を答える。

実装例(C++):

#include<bits/stdc++.h>
#include<atcoder/dsu>

using namespace std;
using namespace atcoder;

int main(){
  int n,q;
  cin >> n >> q;
 
  dsu uf(n+q+5);
  int balls=n;
 
  vector<int> ld(n+1,-1);
  vector<int> inv(n+q+5);
  for(int i=1;i<=n;i++){
    ld[i]=i;
    inv[ld[i]]=i;
  }
 
  for(int tr=0;tr<q;tr++){
    int typ;
    cin >> typ;
    if(typ==1){
      int x,y;
      cin >> x >> y;
      if(ld[y]==-1){continue;}
      else if(ld[x]==-1){
        ld[x]=ld[y];
        inv[ld[x]]=x;
        ld[y]=-1;
      }
      else{
        uf.merge(ld[x],ld[y]);
        ld[x]=uf.leader(ld[x]);
        inv[ld[x]]=x;
        ld[y]=-1;
      }
    }
    if(typ==2){
      int x;
      cin >> x;
      balls++;
      if(ld[x]==-1){
        ld[x]=balls;
        inv[ld[x]]=x;
      }
      else{
        uf.merge(ld[x],balls);
        ld[x]=uf.leader(balls);
        inv[ld[x]]=x;
      }
    }
    if(typ==3){
      int x;
      cin >> x;
      int tg=uf.leader(x);
      cout << inv[tg] << "\n";
    }
  }
  return 0;
}

posted:
last update: