要素列挙可能Union Findの実装例

はじめましての人ははじめまして。[alumi](https://x.com/alumi_02)と申します。
今回は要素列挙可能Union Findの実装例についてまとめました（一応[次回の記事](https://qiita.com/alumite14/items/f4c355720f2a6da88ca5)の伏線）。

## 要素列挙可能Union Findの定義

この記事では、以下の機能のメソッドを持つUnion Findのことを「要素列挙可能Union Find」と呼びます。以下Union Findを初期化する要素数を$N$とします。

- 引数に任意の要素$x \\: (\in N)$を指定
- 要素$x$が属するグループの全要素（$x$含む）をlistなどで返す

本記事ではUnion Find自体の説明は省くので、詳しく知りたい方は以下のリンクをご参照ください。

- [Union Find – 37zigenのHP](https://37zigen.com/union-find/)
- [Union-Find の実装 | アルゴ式](https://algo-method.com/descriptions/136)
- [PythonでのUnion-Find（素集合データ構造）の実装と使い方 | note.nkmk.me](https://note.nkmk.me/python-union-find/)
- [素集合データ構造 - Wikipedia](https://ja.wikipedia.org/wiki/%E7%B4%A0%E9%9B%86%E5%90%88%E3%83%87%E3%83%BC%E3%82%BF%E6%A7%8B%E9%80%A0)
- [Union-Findアルゴリズムと計算科学の話](https://zenn.dev/kaityo256/articles/union_find_physics)
- [Introduction to Disjoint Set (Union-Find Algorithm) - GeeksforGeeks](https://www.geeksforgeeks.org/introduction-to-disjoint-set-data-structure-or-union-find-algorithm/)


## 根の比較による実装 | O(N)

`x`と同じ根を持つ全ての要素を確認する実装です。つまり、

1. `x`の根`rx`を計算
2. 全ての要素の根を`rx`と比較。同じ根をもつ要素を`res`に`append`
3. `res`を`return`

とすれば実装できます。当然全ての要素を確認する必要があるので$O(N)$です。他の実装を変更する必要がないのが利点です。

```diff_python
class UnionFind:
    def __init__(self, n: int) -> None:
        self.par = list(range(n))
        self.sizes = [1] * n

    def _find_root(self, x: int) -> int:
        """Find the root of x with path halving."""
        while x != self.par[x]:
            self.par[x] = self.par[self.par[x]]
            x = self.par[x]
        return x

    def is_same(self, x: int, y: int) -> bool:
        return self._find_root(x) == self._find_root(y)

    def unite(self, x: int, y: int) -> bool:
        rx = self._find_root(x)
        ry = self._find_root(y)
        if rx == ry:
            return False
        if self.sizes[rx] < self.sizes[ry]:
            rx, ry = ry, rx
        self.sizes[rx] += self.sizes[ry]
        self.par[ry] = rx
        return True

+   def get_members(self, x: int) -> list[int]:
+       rx = self._find_root(x)
+       return [i for i, p in enumerate(self.par) if p == rx]
```

`par`の根要素に負の要素数を格納してメモリを削減するような実装の場合、`rx`を先に`res`に入れておかないと`rx`だけ`res`に入らないので注意が必要です。

```diff_python
class UnionFind:
    def __init__(self, n: int) -> None:
+       self.par = [-1] * n

	...

    def get_members(self, x: int) -> list[int]:
        rx = self._find_root(x)
+       return [rx] + [i for i, p in enumerate(self.par) if p == rx]
```

## グループを陽に持つ実装 | O(1)

- 参考：
	- [x.com kemuniku](https://x.com/kemuniku/status/1747668708854755804)
	- [Union-Find木 [いかたこのたこつぼ]](https://ikatakos.com/pot/programming_algorithm/data_structure/union_find_tree#%E3%82%B0%E3%83%AB%E3%83%BC%E3%83%97%E3%81%AE%E8%A6%81%E7%B4%A0%E3%82%82%E4%BF%9D%E6%8C%81)

Union Findの実装を見直し、各要素を実際にsetやlistなどに格納しておき、`unite`の際に親とするグループへ子のグループを実際に併合するようにします。これは実装を見たほうが理解しやすいでしょう。

```diff_python
class UnionFind:
    def __init__(self, n: int) -> None:
+       self.groups = [{i} for i in range(n)]

    def unite(self, x: int, y: int) -> bool:
        s = self.groups[x]
        t = self.groups[y]
        if s is t:
            return False
        if len(s) > len(t):
            s, t = t, s
+       t |= s
+       for i in s:
+           self.groups[i] = t
        return True

    def is_same(self, x: int, y: int) -> bool:
        return self.groups[x] is self.groups[y]
```

この実装では初期化時に`groups`として各要素が一つずつ入った素集合を実際に構成します。そして`unite`時に集合要素が大きいグループへ小さいグループを併合し、小さいグループ内の要素が指すグループの参照先を併合したグループへ振り替えます。
このようにすると、`unite`が$O(\log(N))$となる代わりに`_find_root`が不要となり、`is_same`は2集合$A,B$の比較には$O(min(|A|, |B|))$かかるので$O(N)$となります[^1]。
`x`が属するグループは`group[x]`を呼び出せば得ることができます。

[^1]: `_find_root`を実装し、振替を`_find_root`内で行うようにすればより速いです（[実装例：Union-Find木 [いかたこのたこつぼ]](https://ikatakos.com/pot/programming_algorithm/data_structure/union_find_tree#%E3%82%B0%E3%83%AB%E3%83%BC%E3%83%97%E3%81%AE%E8%A6%81%E7%B4%A0%E3%82%82%E4%BF%9D%E6%8C%81)）。筋の良いハッシュを実装し集合の比較を行ってもいいかも

`is_same`をもっと効率化できないか考えてみます。`roots`を用意し`unite`時の小さいグループの振替を`roots`に対して行うことで、`is_same`は`roots`の比較で済むので$O(1)$となります。`groups`もsetで持つ必要性はないのでlistで持つようにしています（多分ここはどっちでもそんなに計算量変わらないはず）。
`x`が属するグループは`get_members(x)`を呼び出せば得ることができます。

```diff_python
class UnionFind:
    def __init__(self, N: int) -> None:
        self.groups = [[i] for i in range(N)]
+       self.roots = [*range(N)]

    def is_same(self, x: int, y: int) -> bool:
+       return self.roots[x] == self.roots[y]

    def unite(self, x: int, y: int) -> bool:
        rx, ry = self.roots[x], self.roots[y]
        if rx == ry:
            return False
        if len(self.groups[rx]) < len(self.groups[ry]):
            rx, ry = ry, rx
            self.groups[rx].extend(self.groups[ry])
+       for i in self.groups[ry]:
+           self.roots[i] = rx
        self.groups[ry].clear()
        return True

    def get_size(self, x: int) -> int:
        return len(self.groups[self.roots[x]])

    def get_members(self, x: int) -> list[int]:
        return self.groups[self.roots[x]]

```

この構成であれば`x`が属するグループを単なる配列参照$O(1)$で得ることができるので、頻繁にグループを取得したいニーズがある場合に有利です。しかしグループを陽に持つ特性上メモリ使用量が増加します[^2]。またマージテク部分がボトルネックとなり`unite`の計算量が$O(\log(N))$に落ちます。

[^2]: `unite`時に子側のgroupを`clear`すれば、要素数の合計は$N$で抑えられるので、メモリ使用量も$O(N)$で抑えられます。`clear`しない場合でも正常に動作しますが、要素数の合計は最悪$N(N+4)/4$となり、メモリ使用量は$O(N^2)$となります

## グループを陰に持つ実装 | O(group size)

- 参考：
    - [解説 - 第10回 アルゴリズム実技検定 過去問](https://atcoder.jp/contests/past202203-open/editorial/4508)
    - [素集合の要素の列挙と併合 (単方向循環リスト) - noshi91のメモ](https://noshi91.hatenablog.com/entry/2019/07/19/180606)

`unite`、`is_same`の計算量を$O(\alpha(N))$から下げず要素列挙機能を$O(group \ size)$で実現する実装です。陽にグループを持たず、同じグループ間の連結関係を保持します。
具体的には初期化時に`next_v`を用意し、`unite`時に`next_v[rx]`と`next_v[ry]`をスワップし循環リストを構成、実際に列挙が必要なときに`x, next_v[x], next_v[next_v[x]], ...`と辿ることで列挙できるようになります。

```diff_python
class UnionFind:
    def __init__(self, N: int) -> None:
        self.n = N
        self.par = list(range(N))
+       self.next_v = list(range(N))
        self.size = [1] * N

    def _find_root(self, x: int) -> int:
        """Find the root of x with path halving."""
        while x != self.par[x]:
            self.par[x] = self.par[self.par[x]]
            x = self.par[x]
        return x

    def is_same(self, x: int, y: int) -> bool:
        return self._find_root(x) == self._find_root(y)

    def unite(self, x: int, y: int) -> bool:
        rx, ry = self._find_root(x), self._find_root(y)
        if rx == ry:
            return False
+       self.next_v[rx], self.next_v[ry] = self.next_v[ry], self.next_v[rx]
        if self.size[rx] < self.size[ry]:
            rx, ry = ry, rx
        self.size[rx] += self.size[ry]
        self.par[ry] = rx
        return True

+   def get_members(self, x: int) -> list[int]:
+       res = [x]
+       cv = self.next_v[x]
+       while cv != x:
+           res.append(cv)
+           cv = self.next_v[cv]
+       return res
```

このようにすれば、$O(group \ size)$でグループ要素を取得できます。またそれ以外の計算量を悪化させないのが嬉しいです。メモリ使用量的には`next_v`が増えているので、グループを陽に持つ実装とそんなに変わりません。

## 結び

というわけで要素列挙可能Union Findの実装のまとめでした。あまり使う機会はないですが、いざというときに選択肢として持っておきたいですよね（2次元グリッドとかで応用効きそうだし）。改善案や記事の感想など教えてくださると嬉しいです。
良かったらいいね押してね。

## Verification

- [PAST 第10回 H - 連結成分](https://atcoder.jp/contests/past202203-open/tasks/past202203_h)

## References

- Union Findについて
	- [Union Find – 37zigenのHP](https://37zigen.com/union-find/)
	- [Union-Find の実装 | アルゴ式](https://algo-method.com/descriptions/136)
	- [PythonでのUnion-Find（素集合データ構造）の実装と使い方 | note.nkmk.me](https://note.nkmk.me/python-union-find/)
	- [Disjoint-set data structure - Wikipedia](https://en.wikipedia.org/wiki/Disjoint-set_data_structure)
	- [素集合データ構造 - Wikipedia](https://ja.wikipedia.org/wiki/%E7%B4%A0%E9%9B%86%E5%90%88%E3%83%87%E3%83%BC%E3%82%BF%E6%A7%8B%E9%80%A0)
	- [Union-Findアルゴリズムと計算科学の話](https://zenn.dev/kaityo256/articles/union_find_physics)
	- [Introduction to Disjoint Set (Union-Find Algorithm) - GeeksforGeeks](https://www.geeksforgeeks.org/introduction-to-disjoint-set-data-structure-or-union-find-algorithm/)
- グループを陽に持つ実装
	- [x.com kemuniku](https://x.com/kemuniku/status/1747668708854755804)
	- [Union-Find木 [いかたこのたこつぼ]](https://ikatakos.com/pot/programming_algorithm/data_structure/union_find_tree#%E3%82%B0%E3%83%AB%E3%83%BC%E3%83%97%E3%81%AE%E8%A6%81%E7%B4%A0%E3%82%82%E4%BF%9D%E6%8C%81)
- グループを陰に持つ実装
    - [解説 - 第10回 アルゴリズム実技検定 過去問](https://atcoder.jp/contests/past202203-open/editorial/4508)
    - [素集合の要素の列挙と併合 (単方向循環リスト) - noshi91のメモ](https://noshi91.hatenablog.com/entry/2019/07/19/180606)