Union-Find木の解説はこちらの記事などを参照してください。バグや利便性を高められる余地などがあればコメントで教えていただけるとありがたいです。
また、マクロやFOR文の略記を使用しているので、こちらの記事の競プロ用のテンプレートと併せて使用してください。
###修正(2020/08/27)
grouping関数で親ごとで素集合に分割する際に任意の頂点で経路圧縮を行う必要があるのを忘れていたため、バグが発生していました。
また、コード中のコメントアウトを一部修正しました。
###追記(2020/09/24)
要素数をメンバ変数に加えました。
また、素集合系を使い回すためにclear関数を導入しました。
後、気が向いたらtemplate引数に書き直そうと思います。
unionfind.cc
//以下、素集合と木は同じものを表す
class UnionFind{
public:
vector<ll> parent; //parent[i]はiの親
vector<ll> siz; //素集合のサイズを表す配列(1で初期化)
map<ll,vector<ll>> group; //集合ごとに管理する(key:集合の代表元、value:集合の要素の配列)
ll n; //要素数
//コンストラクタ
UnionFind(ll n_):n(n_),parent(n_),siz(n_,1){
//全ての要素の根が自身であるとして初期化
for(ll i=0;i<n;i++){parent[i]=i;}
}
//データxの属する木の根を取得(経路圧縮も行う)
ll root(ll x){
if(parent[x]==x) return x;
return parent[x]=root(parent[x]);//代入式の値は代入した変数の値なので、経路圧縮できる
}
//xとyの木を併合
void unite(ll x,ll y){
ll rx=root(x);//xの根
ll ry=root(y);//yの根
if(rx==ry) return;//同じ木にある時
//小さい集合を大きい集合へと併合(ry→rxへ併合)
if(siz[rx]<siz[ry]) swap(rx,ry);
siz[rx]+=siz[ry];
parent[ry]=rx;//xとyが同じ木にない時はyの根ryをxの根rxにつける
}
//xとyが属する木が同じかを判定
bool same(ll x,ll y){
ll rx=root(x);
ll ry=root(y);
return rx==ry;
}
//xの素集合のサイズを取得
ll size(ll x){
return siz[root(x)];
}
//素集合をそれぞれグループ化
void grouping(){
//経路圧縮を先に行う
REP(i,n)root(i);
//mapで管理する(デフォルト構築を利用)
REP(i,n)group[parent[i]].PB(i);
}
//素集合系を削除して初期化
void clear(){
REP(i,n){parent[i]=i;}
siz=vector<ll>(n,1);
group.clear();
}
};