競技プログラミングで出がちであるデータ構造のUnionFind木のPython3用のコードを書いてみた。競技プログラミングといえばC++の印象は大きいが、最近はPython使用勢が多いため、Pythonで競プロをやっている勢としてサンプルコードなどを公開していけたらと思っている。
まず、AtCoder Typical Contest 001 - B -Union Findの解となるコードを書いてみた。すると このようになる。（外部ページ）UnionFindのクラスは27行目までであって、それを除くと実質コードとしてあるのは11行である。つまり実際のコンテストでUnionFindを使う場面が出現したとき、クラスの部分だけコピー&ペーストすればかなり時間短縮になる。 ただこの例だとfind(x)の使用例までは網羅できていないため、もう一つ例を紹介する。
N頂点M辺からなる単純無向グラフがあり,頂点には1,...,Nと番号がついている.
— しゃかやみ (@shakayami_) 2018年10月26日
k番目の辺はa_kとb_kとを結んでいる。
このとき,1≦i<j≦Nであり,頂点iと頂点jへのパスが存在する頂点の組(i,j)の個数を求めよ.
制約:
N≦10^5,M≦10^5
1≦a_i<b_i≦N(1≦i≦M),i≠jならば(a_i,b_i)≠(a_j,b_j)
Twitterに投稿した何気ない問題であるが、実はこれはUnionFind木を使う典型的な問題である。 UnionFindを初期化してすべての辺を繋げたとする。このとき、「find(x)==find(y)」と「xからyへのルートが存在する」が同値になる。 つまり、find(x)=kとなる頂点集合$X_k$をとってきたとき、$X_k$の任意の異なる2頂点に対してパスが存在する。よって各$k$に対して、${}_{|X_k|}C_2$を総和したものが答えとなる。
少し脱線するが、これは同値類について考えていると解釈することもできる。頂点集合$X=\{1,\ldots,N\}$について、2項関係$\sim$を、$x\sim y\Leftrightarrow$「$x$から$y$へのパスが存在する」とするとこの$\sim$は同値関係を満たす。(つまり、反射律・対称律・推移律を満たす) よって同値関係による商集合$X/\sim$を定めることができる。ここでUnionFind木の特徴として$C(x)\in X/\sim$を適当に取ってきたとき、$\forall y,z\in C(x), \mathrm{find}(y)=\mathrm{find}(z)$が成立する。さらに$\mathrm{find}(x)=\mathrm{find}(y)\Leftrightarrow x\sim y$となるため、find関数を使うことで同値関係にあるかどうかを確かめることができる。 ここで、求める答えは $$\#\{(x,y)\in X^2|x\sim y,x\neq y\}$$ だが、同じ同値類のものだけ考えれば良いので、 $$\sum_{k}\#\{(x,y)\in C(k)^2|x\sim y,x\neq y\}$$ だが、$x,y\in C(k)$なら$x\sim y$なので答えは $$\sum_{k}\#\{(x,y)\in C(k)^2|x\neq y\}$$ としてよいということになる。
最後に使用例ということで例1と例2のコードを載せておく。