サイズ付きUnion Find with size

  • rank付きUFとは違う。rankは木の高さなので
  • サイズはグループの根に保存されるとする
  • size : 木の頂点数
  • 統合ルール
    • sizeの大きい方に統合する
    • sizeが同じ場合は数値の小さい方に統合
// サイズ付きUF
struct UnionFind{
  VI parent;
  VI size; // グループのサイズ(rankとは違う)
  UnionFind(ll sz){
    parent.resize(sz);
    size.resize(sz, 1);
    reset();
  }

  void reset(){
    FOR(i, 0, parent.size()){
      parent[i] = i;
      size[i]=1;
    }
  }

  void unite(ll a, ll b){
    if(a>b) swap(a, b);
    ll rootA = findRoot(a);
    ll rootB = findRoot(b);
    if(rootA==rootB) return;

    if(size[rootA]==size[rootB]){
      // sizeが同じなら、値の小さい方に統合
      if(rootA>rootB) swap(rootA, rootB);
    }
    else{
      if(size[rootA]<size[rootB]) swap(rootA, rootB);
    }
    // rootAを親にするのは共通
    size[rootA] += size[rootB];
    size[rootB] = -1;
    parent[rootB] = rootA;
    return;
  }

  ll getSize(ll a){
    ll r = findRoot(a);
    return size[r];
  }

  ll findRoot(ll a){
    if(parent[a]==a){
      return a;
    }else{
      // 経路圧縮
      return parent[a] = findRoot(parent[a]);
    }
  }

  map<ll, vector<ll> > getGroups(){
    map<ll, vector<ll> > G;
    FOR(i, 0, parent.size()){
      ll r = findRoot(i);
      G[r].push_back(i);
    }
    return G;
  }

  bool is_same_group(ll a, ll b){
    ll rootA = findRoot(a);
    ll rootB = findRoot(b);
    if(rootA==rootB){
      return true;
    }else{
      return false;
    }
  }
};

verified