はじめましての人ははじめまして。alumiと申します。
今回は要素列挙可能Union Findの実装例についてまとめました（一応次回の記事の伏線）。

要素列挙可能Union Findの定義

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

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

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

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

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

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

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

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に入らないので注意が必要です。

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
- Union-Find木 [いかたこのたこつぼ]

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

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)$となります¹。
xが属するグループはgroup[x]を呼び出せば得ることができます。

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

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)$で得ることができるので、頻繁にグループを取得したいニーズがある場合に有利です。しかしグループを陽に持つ特性上メモリ使用量が増加します²。またマージテク部分がボトルネックとなりuniteの計算量が$O(\log(N))$に落ちます。

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

参考：
- 解説 - 第10回アルゴリズム実技検定過去問
- 素集合の要素の列挙と併合 (単方向循環リスト) - noshi91のメモ

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]], ...と辿ることで列挙できるようになります。

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 - 連結成分

References

Union Findについて
グループを陽に持つ実装
- x.com kemuniku
- Union-Find木 [いかたこのたこつぼ]
グループを陰に持つ実装
- 解説 - 第10回アルゴリズム実技検定過去問
- 素集合の要素の列挙と併合 (単方向循環リスト) - noshi91のメモ

_find_rootを実装し、振替を_find_root内で行うようにすればより速いです（実装例：Union-Find木 [いかたこのたこつぼ]）。筋の良いハッシュを実装し集合の比較を行ってもいいかも ↩
unite時に子側のgroupをclearすれば、要素数の合計は$N$で抑えられるので、メモリ使用量も$O(N)$で抑えられます。clearしない場合でも正常に動作しますが、要素数の合計は最悪$N(N+4)/4$となり、メモリ使用量は$O(N^2)$となります ↩