networkXでUnionFind

atcoderで頻出なunionfindを、自前ライブラリ実装ではなく、networkXで実装する方法の紹介です。
networkXについて調べていたら、unionfindを発見したが、networkXのunionfindの日本語記事が見つからなかったので紹介します。

以下のコードは現在のatcoderのpythonのバージョンで実行やACをしたものです。
Python: 3.8.2
NetworkX: 2.4

テンプレート

networkxuf.py

from networkx.utils import UnionFind
uf = UnionFind(range(1,6))
uf.union(1, 2)  # 1と2をunion
uf.union('atcoder', (3, 'chokudai')) # 'atcoder'と(3,'chokudai')をunion
uf.union(2, 3, 4)
print(uf[2] == uf[(3, 'chokudai')])  # 2と(3,'chokudai')が同じグループか判定(uf[a]はaの根を返す)
# False
print(uf['atcoder'] == uf[(3, 'chokudai')])
# True
print(uf.weights[uf[5]])  # 5が属する集合の大きさを返す
# 1
print(uf.weights[uf[4]])
# 4
for group in uf.to_sets():  # すべてのグループのリストを返す
    print(group)
"""
{1, 2, 3, 4}
{'atcoder', (3, 'chokudai')}
{5}
"""
for item in uf:
    print(item)
"""
1
2
atcoder
(3, 'chokudai')
3
4
5
"""

使い方説明

from networkx.utils import UnionFind
uf = UnionFind()
unionfindをimportしてインスタンスを生成します。よく見る自前実装の方式と違う特徴は、インスタンス生成時に要素数を指定しません。代わりに、要素のリストを渡します。また、要素のリストを渡さなくても、後述のメソッドで要素を参照するときに、要素が無ければ自動で生成されます。

uf.union(1, 2)
uf.union('atcoder', (3, 'chokudai'))
uf.union(2, 3, 4)
要素を結合するメソッドです。よく見る自前実装と違う特徴は、hashableなら何でも要素に出来ることです。文字列でもタプルでも要素に出来ます。また、3個以上の要素をまとめて結合することも
できます。

print(uf[2] == uf[(3, 'chokudai')])
print(uf['atcoder'] == uf[(3, 'chokudai')])
要素が同じグループか判定するメソッドが見つからなかったので、代わりにこう書きます。uf[x]がxの根を返すので、根が同じかどうかで判定しています。

print(uf.weights[uf[5]])
print(uf.weights[uf[4]])
uf.weights[根]でグループの大きさが取得できます。

for group in uf.to_sets():
print(group)
for item in uf:
print(item)
uf.to_sets()でグループごとにsetになったイテレータを返します。
ufで個別の要素のイテレータを返します。

実践例

AtCoder Library Practice Contest A Disjoint Set Union　

ACL Beginner Contest C Connect Cities(コンテスト本番での使用)

AtCoder Beginner Contest 177 D Friends

エイシング プログラミング コンテスト 2019 C Alternating Path

https://atcoder.jp/contests/aising2019/submissions/17055243
数値以外もunionfindできるのでこうすることもできます。

AtCoder Beginner Contest 040 D 道路の老朽化対策について(ギリギリ！)

CODE FESTIVAL 2016 Final C Interpretation

https://atcoder.jp/contests/cf16-final-open/submissions/17078026
3つ以上をを同時にunionfindできるので、こういうふうに簡潔に書くこともできます

AtCoder Regular Contest 027　B 大事な数なのでZ回書きまLた。

https://atcoder.jp/contests/arc027/submissions/17078209
文字列もunionfindできるのでこういうふうに実装できます。

速度は？

networkXなので当然遅いです。
以下が上記実践例と、unionfind自前実装との速度の比較です。

問題	自前実装	networkX実装
Disjoint Set Union	401ms	924ms
Connect Cities	211ms	843ms
Friends	619ms	1374ms
Alternating Path	486ms	1373ms
道路の老朽化対策について	1158ms	1854ms(右の実装を高速化してTLE回避)

このように、自前実装と比べて2倍と200ms程度遅くなります。しかし、簡易な軽い問題なら、十分間に合います。簡易な問題なら、実行時間だけでなく、実装時間も大事ですので、ライブラリを探してコピペすることなく、手軽に実装できるnetworkXは十分選択肢に挙がるでしょう。

まとめ

networkXにはまだまだ知らない機能がある。やはりライブラリの豊富さはpythonの強み。（multiset以外...）