グラフ系アルゴリズム:Union-Find 木

先日書いた記事の

katoufujibanana.hatenadiary.jp

にて、グラフ上のある2点が接続しているかどうかの判定アルゴリズムとして、Union-Find の名前を出しましたが、どのようなアルゴリズムだったかは触れなかったので今回は Union-Find を取り上げます。

まずは、Union-Find を使わず、深さ優先探索幅優先探索で2点の接続判定を行う場合、最小計算量  O(1)、最悪計算量  O(n) になります。

Union-Find を使う場合、最小計算量  O(1)、最悪計算量  O(log(n)) となります。さらに、Union-Find は接続判定を行えば行うほど、計算量  O(1) で計算できる点の組が増えていきます。すごい。

その代償として連結しているか否かの判定に特化するために、グラフ構造をすごい勢いで破壊していきます。元のグラフ構造に基づいた処理もしたい場合は Union-Find 用のデータと別にグラフを保持するデータを持っておくようにしてください。

以下スクリプトです。

class UnionFindTree:
    def __init__(self, n):
        self._n = n
        self.parent = [-1] * n

    def merge(self, a, b):
        x = self.leader(a)
        y = self.leader(b)

        if x == y:
            return x

        if -self.parent[x] < -self.parent[y]:
            x, y = y, x

        self.parent[x] += self.parent[y]
        self.parent[y] = x

        return x

    def same(self, a, b):
        return self.leader(a) == self.leader(b)

    def leader(self, a):
        path = []
        for _ in range(self._n):
            parent = self.parent[a]
            if parent < 0:
                break
            path.append(a)
            a = parent

        for p in path:
            self.parent[p] = a

        return a

では、スクリプトを1行ずつ見ていきましょう。

class UnionFindTree:
    def __init__(self, n):
        self._n = n
        self.parent = [-1] * n

コンストラクタです。n はグラフのノード数です。

Union-find では、グラフの正確な構造は保持せず、あるノードが接続する親のノードのみ記憶します。これを記憶する配列が self.parent です。初期値は -1 であり、これはどのノードも親を持たないことを表します。

self.parent は、接続する親のノードだけでなく、自分が接続しているグラフのサイズも表現します。値が正の値である場合は自分の接続先の親ノード、負の値である場合は自分が接続しているグラフのサイズを表現します。後述の merge メソッドや、leader メソッドで self.parent を更新しますが、このルールが守られるように更新されていることに注意してください。

    def merge(self, a, b):
        x = self.leader(a)
        y = self.leader(b)

        if x == y:
            return x

        if -self.parent[x] < -self.parent[y]:
            x, y = y, x

        self.parent[x] += self.parent[y]
        self.parent[y] = x

        return x

merge メソッドの説明に入ります。leader メソッドの詳細は後に回しますが、とりあえず、「a からグラフを親方向に辿って行ったときの一番上のノード」を返すメソッドであると理解いただければ大丈夫です。

merge メソッドを実行するのは、ノード a とノード b の間に辺が追加されたときです。merge メソッドによって Union-find のデータ構造を更新します。

        if x == y:
            return x

すでに、a と b が同じ leader を持つ場合、データ構造は更新しません。

        if -self.parent[x] < -self.parent[y]:
            x, y = y, x

        self.parent[x] += self.parent[y]
        self.parent[y] = x

そうでない場合、「x のサイズに、y のサイズを可算」し、「y の親を x」にします。

前段の処理として、サイズが大きい方にサイズが小さい方を接続するようにすると、leader メソッドの計算量が抑えられるので、お互いのサイズを比較して大きい方に接続できるように必要に応じて x と y を入れ替えておきます。

    def same(self, a, b):
        return self.leader(a) == self.leader(b)

same メソッドは、a と b が接続しているか否かを判定します。

    def leader(self, a):
        path = []
        for _ in range(self._n):
            parent = self.parent[a]
            if parent < 0:
                break
            path.append(a)
            a = parent

        for p in path:
            self.parent[p] = a

        return a

leader メソッドです。

    def leader(self, a):
        path = []
        for _ in range(self._n):
            parent = self.parent[a]
            if parent < 0:
                break
            path.append(a)
            a = parent

この for ループで、a を起点として親方向にノードを辿っていき、parent の値が負になったとき、つまり一番上のノードにたどり着いた際に break し、ループを抜け出します。値の更新の流れから、ループを抜け出したときの a が一番上のノードを表します。

ループの中で、path という配列に、辿っていくノードを保存しています。

        for p in path:
            self.parent[p] = a

        return a

一番上のノードが確定した後は、path に格納されているノード(つまり、今までに辿ってきたノード)のすべての親を、一番上のノードに設定しなおします。
この工夫により、一番上のノードとの距離の小さくなるノードが増えてゆき、leader メソッドの計算量をどんどん抑えていくことができるようになっています。

具体的に計算量がどれくらい抑えられているのかは、またいつか別の記事で。