はじめに
本記事では，UnionFindクラスが必要となる，無向グラフに関する問題（AtCoder ABCコンテスト 292のD問題）を解説しています．本記事で解説する問題を十分に理解し，解けるようになることで，UnionFindの仕様方法・仕組みだけでなく，クラスやオブジェクト指向プログラミングへの慣れが期待できます．
1. UnionFindについて
まず，UnionFindクラスについて，詳細は以下のサイトが非常に参考になります．実装の更なる効率化についても言及されていますので，今後もお世話になるでしょう．
PythonでのUnion-Find（素集合データ構造）の実装と使い方 | note.nkmk.me
簡単にまとめると
UnionFind（クラス）とは，木構造（グラフ）に新たな要素を追加するunionメソッドと，指定した要素が，どの要素を根とした集合（連結成分；次項で解説）に属するかを判定するfindメソッドという2つの機能を中心とした，データ構造です．
クラスとは，データの値だけでなく，様々な機能（メソッド）も兼ね備えたデータ構造ですが，UnionFindクラスの持つメソッドには，例えば以下のような種類があります．
same(x, y) : 指定した2つの要素x, yが，同じグループに属するかを判定するメソッド
size(x) : 指定した要素xの属するグループ（木）の要素数を返すメソッド
代表的なメソッドは他にもいくつかあり，また，解くべきに応じて，必要なメソッドを追加実装することも可能です．しかし，実装がやや複雑なため，慣れないうちは1つの典型的なUnionFindクラスを使い回すのが良いです．ひとつの典型例を時間をかけて理解し，手を動かして使用していれば，おのずと全体像が見えてきます．
とはいえ，UnionFindクラスについて，Pythonによる実装例はC++のそれ程は充実していません．
そこで，私がネットで調べた限りで最も分かりやすく，かつ汎用的な実装をされていた上記の記事を参考に，UnionFindクラスを記述しました．多くの問題に応用できる機能（メソッド）を備えた，UnionFindクラスの実装例を紹介します．
実装例
from collections import defaultdict
class UnionFind():
"""
Union Find木クラス
Attributes
--------------------
n : int
要素数
patents : list
指定した要素の親（1つ上の）要素を格納
指定した要素が根の場合は，
-(グループの要素数) を格納
=> sizeメソッドに反映
"""
def __init__(self, n):
self.n = n
self.parents = [-1] * n
def find(self, x):
"""
ノードxの根を見つける
"""
if self.parents[x] < 0:
return x
else:
self.parents[x] = self.find(self.parents[x])
return self.parents[x]
def union(self, x, y):
"""
木に新たな要素を併合（マージ）
Parameters
---------------------
x, y : int
併合するノード
"""
x = self.find(x)
y = self.find(y)
if x == y:
return
if self.parents[x] > self.parents[y]:
x, y = y, x
self.parents[x] += self.parents[y]
self.parents[y] = x
def size(self, x):
"""
xの属する木のサイズ
"""
return -self.parents[self.find(x)]
def same(self, x, y):
"""
x, yが同じ木に属するか判定
"""
return self.find(x) == self.find(y)
def members(self, x):
"""
xの属する木に属する要素をリストで返す
"""
root = self.find(x)
return [i for i in range(self.n) if self.find(i) == root]
def roots(self):
"""
全ての根をリストで返す
"""
return [i for i, x in enumerate(self.parents) if x < 0]
def group_count(self):
"""
グループの数を返す
"""
return len(self.roots())
def all_group_members(self):
"""
全てのグループの要素情報を辞書で返す
"""
group_members = defaultdict(list)
for member in range(self.n):
group_members[self.find(member)].append(member)
return group_members
def __str__(self):
"""
print(uf)で全てのグループの要素情報を簡単に出力する
"""
return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())2. AtCoder ABC270 D – Unicyclic Components
N個の頂点（ノード）とM個の辺（エッジ）からなる無向グラフは，1つ以上の連結成分からなります．連結成分とは，下図のように，独立した各々の部分グラフのことです．下図では，N = 9，M = 1 + 3 + 5 = 9の場合を示します．
本問は，各連結成分について，含まれる要素の数と辺の数が等しいかを調べ，すべての連結成分で「要素数=辺の数」が成り立つかを調べる問題です．上の例の場合，左（要素2個，辺1本）と右（要素4個，辺5本）は要件を満たさず，真ん中（要素3個，辺3本）は要件を満たします．
問題文の詳細はこちら：D – Unicyclic Components (atcoder.jp)
3. UnionFindクラスを用いた解法
上の問題を解くポイントは
- 各連結成分に含まれる要素数（連結成分の大きさ）を数える．
- 各連結成分に含まれる辺の数を数える．
という，一見単純な作業を行って，両者の結果を比較することです．前者については，UnionFindクラスの標準的なメソッドである，sizeに，連結成分の根の要素を引数として渡せばよいです．
size(x = 要素) = (要素xの所属する連結成分のサイズ（頂点の数）)一方，後者については，欲しい値そのものを返してくれるメソッドがありません．しかし，全ての要素情報をUnionFind木にインプット（union）した後，各辺がどの根の連結成分（グループ）に属するかを調べれば，各々の連結成分（グループ）の持つ情報に分けることができます．
注意点：
木構造ではなく，本問題で扱うような閉路を含むようなグラフにおいて，根（に相当するノード）は，最も値の小さいノードとなる．
実装上では，各辺を構成する2つの要素x, yの内のどちらかについて，findメソッドによって，所属する連結成分の根を調べればよいです．
実装例
n, m = map(int, input().split())
# 辺（エッジの全データを格納）
ed = []
# インスタンスを生成
uf = UnionFind(n)
for _ in range(m):
# 0-indexed
x, y = map(int, input().split())
x, y = x - 1, y - 1
# 新たなノードとエッジを追加
uf.union(x, y)
ed.append((x, y))
# 根をrとする連結成分（グループ）の辺の数を格納
es = {r: 0 for r in uf.roots()}
for x, y in ed:
# ノードx（yでも良い）の根を見つける
r = uf.find(x)
# 根rのグループに辺の数を+1
es[r] += 1
cond = all(uf.size(l) == es[l] for l in uf.roots())
print("Yes" if cond else "No")実行例
上で3つのグループを図示した入力の場合（インデックスの始まりは1）．
9 9
1 2
3 4
3 6
6 4
5 7
5 9
7 9
9 8
8 7
>> No最後に
タイトルでは「連結成分について分割」と称しましたが，実際にグラフの隣接リスト（辺の全情報）を，連結成分ごとに分割したわけではありません．しかし，一度UnionFindクラスからインスタンスを生成した後，各辺の根をfindメソッドで照合することによって，連結成分毎に辺の数をカウントできることが分かりました．この方法は，連結成分に係る様々な問題で応用できるのではないかと思います．
コメント