はじめましての人ははじめまして。alumiと申します。
昨日（2024-09-07）のABC370のD問題の模範回答でUnionFindの面白い拡張の例示がなされていたので、他にどんなことができそうか考察してみました。

ABC370 D Cross Explosion

問題の概要は以下のような感じです（正確な内容は問題ページを参照ください）。

ボンバーマンのステージのような、最初全てが壁でできた$H \times W$のグリッドが与えられる
クエリとして$1 \le r \le H$, $1 \le c \le W$なる$r,c$が$Q$回与えられる
- グリッドの位置$(r,c)$が壁なら、$(r,c)$の壁を一つだけ破壊する
- グリッドの位置$(r,c)$が壁でないなら、$(r,c)$から上下左右4方向それぞれで、最初に到達する壁を破壊する（1クエリで最大4つ破壊）

制約から愚直に最初に到達する壁を探索すると間に合わないため、クエリとして与えられた$(r,c)$からみて、最初に到達する壁の位置を効率的に取得する必要があります。

解答としてはSortedSetによる方法（壁があるなら消し、ないならlt, gtで両サイドを取得。コード例）やセグメント木上の二分探索（01とorを載せて、min_left, max_rightで$x \le 0$を満たすindexを探索。コード例）がありますが、模範解答としてUnion Findを拡張した方法がありました。これは次のような方針に基づく解法です。

各行ごと・各列ごとにunion findのインスタンスを用意
現在位置が壁かどうかの情報も用意
各クエリに回答
1. もし$(r,c)$が壁なら、$(r,c)$を壁でなくし、四方向の壁でないマスとunite
2. もし$(r,c)$が壁でないなら、$(r,c)$が属する行素集合の左端・右端、列素集合の上端・下端を取得し、それぞれで四方向の壁でないマスとunite

この太字の機能を実現するためには、Union Findにたった7行追加するだけでOKです。すごい！

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

    def _find_root(self, x: int) -> int:
        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
        if self.size[rx] < self.size[ry]:
            rx, ry = ry, rx
+       self.L[rx] = min(self.L[rx], self.L[ry])
+       self.R[rx] = max(self.R[rx], self.R[ry])
        self.size[rx] += self.size[ry]
        self.par[ry] = rx
        return True

+   def get_LR(self, x: int) -> tuple[int, int]:
+       rx = self._find_root(x)
+       return self.L[rx], self.R[rx]

つまりmerge（unite）の際にそれぞれの素集合の代表値を比較し、小さい方・大きい方を新たな代表値として採用するという操作を追加すればいいです。簡単～。

これにより自分が属する素集合の最大値・最小値を取得できるので、そのデータを使って上記の方針でuniteすればいいですね（コード例）

応用と拡張性

前述の問題ではグリッドに対する要素を持っていましたが、この拡張はそれ以外にも次のようなシチュエーションで使えそうです

要素数$N$の整数列$A$に対し、要素$i,j$が属する素集合同士をmergeするクエリと、要素$i$が属する素集合の要素の総和・最大値・最小値を取得するクエリに答える

これは総和・最大値・最小値保持用の配列を$A$で初期化し、unite時に比較すればいいですね。他にも二項演算なら大体載せられそうです（使い道はわからんけど）。

class UnionFind:
    def __init__(self, N: int, A: list) -> None:
        ...
        self.sum = A[:]
        self.max = A[:]
        self.min = A[:]
    ...

    def unite(self, x: int, y: int) -> bool:
        """Unite the sets containing x and y."""
        ...
        self.sum[rx] = self.sum[rx] + self.sum[ry]
        self.max[rx] = max(self.max[rx], self.max[ry])
        self.min[rx] = min(self.min[rx], self.min[ry])
        ...

他の応用が思いついたら追記します。

任意のモノイド積に対応するヘルパークラス

UnionFind自体の実装を汚したくない場合、下記のようにUnionFind classを継承したヘルパークラス内に上記の計算を実装すると簡潔に記載できます。

class CommutativeMonoidUnionFind(UnionFind):
    def __init__(self, N, init_arr, op):
        super().__init__(N)
        self.total_product = init_arr
        self.op = op

    def unite(self, x: int, y: int) -> bool:
        """Unite the sets containing x and y."""
        rx, ry = self._find_root(x), self._find_root(y)
        if rx == ry:
            return False
        new_product = self.op(self.total_product[rx], self.total_product[ry])
        super().unite(rx, ry)
        new_rx = super()._find_root(rx)
        self.total_product[new_rx] = new_product
        return True

    def get_total_product(self, x):
        return self.total_product[self._find_root(x)]

N: int = int(input())
ds = CommutativeMonoidUnionFind(N, list(range(N)), min)

結び

というわけでUnion Findの拡張に関する記事でした。解説で知って感動したので、いろんな人とこの感動を共有したくて勢いのまま書きました。乱文ご容赦ください。

Verification