PythonのNetworkXのUnionFind(DSU)の特徴のメモ

# 最初に

AtCoderのPythonで使える`networkx`の`UnionFind`のメモです。
AtCoderで使えるけれど、あまり速くない気がします。

`networkx`の`UnionFind`をRubyに移植したい機会があり、
実装を見る機会があったので、特徴をまとめたいと思います。

ただ、普段Pythonを使わず、Pythonの用語じゃないかもしれないので、違ってたらコメントでも修正でも書いてもらえたら、嬉しいです。

# UnionFindの特徴

DSU(disjoint set union)とも言いますね。

出来ることは主に
- 要素の結合して、同じグループとする。
- それぞれの要素が同じグループにいるかどうかを素早く確かめる。

ですね。

# 主なメソッドの特徴

主なメソッドの特徴を書いていきます。

## `__init__`

インスタンスを作るとき、初期化はリストや範囲型の変数を引数で渡します。
`for`文で回せるイテラブルなものは引数にとれます。
何も引数を渡さなければ、空リストと同じ扱いです。
引数にとるものは、グラフの要素(node)を表すリストや範囲型です。

内部では、`parents`と`weights`という辞書型のインスタンス変数を持ちます。
- `parensts`: それぞれの要素が持つ親を指す辞書。
- `weights`: それぞれの要素が持つ自身と子の要素数を指す辞書。

競プロの多くの本屋記事では配列でデータを格納するので、辞書型なのは珍しいですね。
1配列でも十分なケースが多いのですが、 NetworkXのUnionFindは2つの辞書型で機能しているので、メモリをたくさん悔い、単純な問題では最適化できておらずスピードの面で若干性能が悪そうな気がします。

インスタンスを初期化した際は、要素は全て独立なので、
- `parents`の全要素: それぞれ自分自身を親とする。
- `weigths`の全要素: サイズ1

として初期化されます。

```python
import networkx.utils as nxu
uf = nxu.UnionFind(range(3))

print(uf.parents) # {0: 0, 1: 1, 2: 2}
print(uf.weights) # {0: 1, 1: 1, 2: 1}
```

```python
import networkx.utils as nxu
uf = nxu.UnionFind(['a', 'b'])

print(uf.parents) # {'a': 'a', 'b': 'b'}
print(uf.weights) # {'a':  1,  'b':  1}
```

```python
import networkx.utils as nxu
uf = nxu.UnionFind()

print(uf.parents) # {}
print(uf.weights) # {}
```

## `union`と`[]`(`__getitem__`)

`union`で結合。
`[]`でグループの代表を返します。
そのため、グループの代表どうしを比べて、一緒なら同じグループとします。

なお、直接、要素どうしが同じグループか調べるメソッドは実装を見る限りありませんでした。
そのため、`uf[x] == uf[y]`の形で判定します。

```python
import networkx.utils as nxu
uf = nxu.UnionFind(range(10, 20))

# 最初、それぞれの代表は自分自身であることを確かめる
print(uf[10], uf[15]) # 10 15

# 要素10と要素15が同じグループか確かめる
print(uf[10] == uf[15]) # False

# 要素10, 14を連結させて同じグループにする
uf.union(10, 14)

# 要素12~15を連結させて、同じグループにする
uf.union(12, 13, 14, 15)

# 要素10と要素15が同じグループか確かめる
print(uf[10], uf[15]) # 10 10
print(uf[10] == uf[15]) # True
```

`union`は、可変長引数である点が面白いところです。  
このせいで、2要素の結合に最適化できておらず、遅い感じがします。

面白い点で注意点でもあるのが、
初期化時に引数として登録しなかった要素があっても、
`union`や`[]`で引数にとると、エラーにせず
それもUninFind木の新しい要素として迎えるところですね。

```python
import networkx.utils as nxu

uf = nxu.UnionFind()

uf[16]
uf.union(13, 12, 11)

print(uf[13], uf[11]) # 13 13
```

間違った引数を渡さないことに注意です。

## `__iter__`

`__iter__`を定義しており、 
`iter(self.parents)`を返すので、
それぞれの要素を`for`文で回すことが出来ます。

```python
import networkx.utils as nxu

uf = nxu.UnionFind(['a', 'b'])
for e in uf:
  print(e)
```

## `to_sets`

それぞれの要素を同じグループごとにまとめる`generator`を返します。

```python
import networkx.utils as nxu

uf = nxu.UnionFind("xyz")

sets = sorted(map(sorted, uf.to_sets()))
print(sets) # [['x'], ['y'], ['z']]

uf.union('x', 'y')
sets = sorted(map(sorted, uf.to_sets()))
print(sets) # [['x', 'y'], ['z']]

print(type(uf.to_sets())) # <class 'generator'>
```

# その他

グループのサイズなどを求めるときは、次のようにします。

```python
import networkx.utils as nxu

uf = nxu.UnionFind("xyz")

uf.union('x', 'y')
print(uf.weights['y'])     # 1 自身と子のサイズ
print(uf.weights[uf['y']]) # 2 自身の属するグループのサイズ
```

# 追記: 競技プログラミング AtCoderで使って解いてみた

- [ATC 846ms Python 3.8.2](https://atcoder.jp/contests/atc001/submissions/36295204)
  【参考】Rubyのac-library-rbのUnionFindは321msですので、特別速くはないと思います。

2022/11/13のAtCoderのABC277を見てたらUnionFindで解ける問題が出題されてたので解いてみました。
- [AC 1881ms: 最小限にコードの綺麗さ・わかりやすさ優先で書いたコード](https://atcoder.jp/contests/abc277/submissions/36456627)
- [AC 1553ms: setを使って、重複を減らしループ回数等を減らした](https://atcoder.jp/contests/abc277/submissions/36456948)

AC(正解)は出来ましたが、実行制限時間が2秒のところ1881ms~1553msなのでギリギリです。
まぁ、競技プログラミング的にはもっと性能を改善してもらいたいところだと思います。



## 参考

[networkx.utils.union_find — NetworkX 2.8.8 documentation](https://networkx.org/documentation/stable/_modules/networkx/utils/union_find.html)