ABC183 F - Confluence

問題

atcoder.jp

解法

生徒の合流をUnionFind木で管理する。

各生徒の集合状態をmapで管理し、生徒たちが合流したときにmapをマージする。 このとき、mapの小さいほうから大きいほうへマージすることで、マージに必要な計算量を抑えることができる。

「データ構造をマージする一般的なテク」と呼ばれるらしい。 証明はこちら記事が参考になりました。 データ構造をマージする一般的なテク

コード

struct UnionFind {
  vector< int > data;
  vector<map<int,int>> mp;
 
  UnionFind(int sz) {
    data.assign(sz, -1);
    mp.assign(sz, map<int,int>());
  }
  int root(int x) {
        if (data[x] < 0) return x;
        else return data[x] = root(data[x]);
  }
  bool same(int x, int y) {
    return root(x) == root(y);
  }
  bool unite(int x, int y) {
    x = find(x), y = find(y);
    if(x == y) return (false);
    if(data[x] > data[y]) swap(x, y);
    for(auto it : mp[y]){
      mp[x][it.fi]+=it.se;
    }
    mp[y] = map<int,int>();
    data[x] += data[y];
    data[y] = x;
    return (true);
  }
 
  int find(int k) {
    if(data[k] < 0) return (k);
    return (data[k] = find(data[k]));
  }
 
  int size(int k) {
    return (-data[find(k)]);
  }
};

int main() {
  int n,q;
  cin >> n >> q;
  vi C(n);
  UnionFind tree(n);
  rep(i,n){
    cin >> C[i];
    C[i]--;
    tree.mp[i][C[i]]=1;
  }
  while(q--){
    int type,a,b;
    cin >> type >> a >> b;
    a--;b--;
    if(type == 1) tree.unite(a,b);
    else cout<<tree.mp[tree.find(a)][b]<<endl;
  }
  return 0;
}