Union_Find (Disjoint Set Union) クラス for Python3

Share

とりあえずそこそこ形になったので、コードを。一応実コンテストでも使用してるので、そんなに間違いはないはず……。

2020/11/09 適当な拡張クラスを追加。グリッドのつながりをUnion Findで扱いたい時などに使えるかもしれない。

# Union_Findクラスfor Python3
# 書いた人: scrblbug
# サイトURL: https://miaoued.net Twitter: @scrblbug
# なんにせよ分かりやすさ重視で……コメントは過剰につけています
# Union_Find(N):要素数NのUnion_Find木を作成
# .parent:親要素管理リスト
# .rank:木の高さ管理リスト
# .group_count:現在のグループ数
# .N:全体の要素数
# .find(x):最上位の親(グループリーダー)を取得
# .unite(x, y):xとyのグループを統合しつつ、(統合先リーダー, 統合元リーダー)or統合不要なら(-1, -1)を返す
# .samep(x, y):xとyが同じグループかどうかを判定
# .get_group_count(x):xの所属するグループのメンバー数を取得
# .get_group_members(x):xの所属するグループのメンバーをリストで取得
# .get_all_group_count():全てのリーダー:グループメンバー数を辞書形式で取得
# .get_all_group_members():全てのリーダー:[メンバー一覧]を辞書形式で取得

class Union_Find:
    # 親管理リストと高さ管理リストを初期化し、
    # 要素N個のUnion-Find森を作成する。
    # 親管理リストは、基本的には自分のひとつ上の親を表すが、
    # 値が負の場合には、自身が最上位の親(リーダー)であることを表し、
    # 自分を含めたグループの人数を管理することとする
    def __init__(self, N):
        self.parent = [-1] * N
        self.rank = [0] * N
        self.group_count = N
        self.N = N

    # xの所属するグループのリーダーを返す
    def find(self, x):
        # 自分自身がリーダーなら、自分を返す
        if self.parent[x] < 0:
            return x

        # 再帰的に捜索し、見つかれば繋ぎ変えておく
        # (計算量が増える=面倒くさいので)高さ管理は行わない
        par = self.find(self.parent[x])
        self.parent[x] = par
        return par

    # xとyのグループを統合する
    def unite(self, x, y):
        # それぞれのリーダーに対する操作を行うことになる
        x = self.find(x)
        y = self.find(y)

        # リーダーが同じなら何もする必要がない
        if x == y:
            return (-1, -1)

        # 木の高さが同じ場合:
        # グループの人数を合計しつつ適当に繋ぎ、繋げられた方の高さを1増やす
        if self.rank[x] == self.rank[y]:
            self.parent[x] += self.parent[y]
            self.parent[y] = x
            self.rank[x] += 1
            self.group_count -= 1
            return (x, y)

        # 木の高さが違うなら、低い方を高い方につなぐ
        if self.rank[x] < self.rank[y]:
            x, y = y, x
        self.parent[x] += self.parent[y]
        self.parent[y] = x
        self.group_count -= 1
        return (x, y)

    # xとyが同じグループかどうかを調べる
    def samep(self, x, y):
        return self.find(x) == self.find(y)

    # xの所属するグループのメンバー数を返す
    def get_group_count(self, x):
        x = self.find(x)
        return -self.parent[x]

    # xの所属するグループのメンバーをリストで返す
    def get_group_members(self, x):
        x = self.find(x)
        return [i for i in range(self.N) if self.find(i) == x]

    # 全ての{リーダー:グループメンバー数}を辞書形式で返す
    def get_all_group_count(self):
        return {idx:-n for idx, n in enumerate(self.parent) if n < 0}

    # 全ての{リーダー:[メンバー一覧]}を辞書形式で返す
    # もしもこれが欲しいだけなら、グラフ探索するほうが速いんじゃないかな……
    def get_all_group_members(self):
        agm_dic = {}
        for i in range(self.N):
            agm_dic.setdefault(self.find(i), [])
            agm_dic[self.find(i)].append(i)
        return agm_dic

# 以下はimmutableなオブジェクトをUnion_Findで扱う拡張クラス
# 適当にでっち上げたのでコメント無し&不具合が多いかも……
class Union_Find_Objects(Union_Find):
    def __init__(self):
        super().__init__(0)
        self.obj_to_idx = {}
        self.idx_to_obj = []
    
    def obj_add(self, obj):
        self.parent.append(-1)
        self.rank.append(0)
        self.idx_to_obj.append(obj)
        self.obj_to_idx[obj] = self.N
        self.N += 1
        self.group_count += 1
    
    def obj_find(self, obj):
        x = self.obj_to_idx[obj]
        super().find(x)
        return self.idx_to_obj[x]
    
    def obj_unite(self, obj_x, obj_y):
        x = self.obj_to_idx[obj_x]
        y = self.obj_to_idx[obj_y]
        x, y = super().unite(x, y)
        if x != -1:
            return (self.idx_to_obj[x], self.idx_to_obj[y])
        else:
            return (-1, -1)
    
    def obj_samep(self, obj_x, obj_y):
        x = self.obj_to_idx[obj_x]
        y = self.obj_to_idx[obj_y]
        return super().samep(x, y)

詳細については、下のスライドが詳しい。

以下、グラフとか頂点とか辺とかを使わないで平易に説明してみる。

砕けた言い方をすると、友達の友達は友達、という関係性の判定を簡単に行うためのものだ。AとB、BとC、CとD、DとEが友人であったとき、AとEが友人かどうか判定する際、通常の探索ならばA→B→C→D→Eとたどる必要がある。5人ならどうということは無いが、これが100人、1000人のグループが100個、1000個となるとその度に探索をするのは大変(=計算量が大)だ。

そこで、まず各グループに上司→部下、あるいは連絡網のような階層構造を設け、各人がどのグループにいるのかを簡単に問い合わせできるようにする(階層構造なので、計算量はO(log N)となる)。グループにA~Zの人がいるとしても、Zの上司はM、Mの上司はE、Eの上司はAで、Aはグループリーダー、といった感じで探索を進めようというイメージだ。これにより、対象が誰をリーダーとするグループにいるのかを簡単に(=少ない計算で)判定できる(階層構造なので、O(log N)となる)。

このようにした時、xとyが友達かどうか、つまり同じグループにいるかどうかを見るためには、要するにxのグループのリーダーとyのグループのリーダーが一致しているかどうかを見ればよいことになる。これなら簡単。やったぜ。

……だったら、最初っからグループの箱をつくって、その中にメンバーを放り込めばいい(あるいはグループのメンバーに名札をつければいい)んじゃない?と思うかもしれないが、そうすると、異なるグループのメンバー同士が友達になった時に困る。なぜなら、「その度に」どちらか片方のグループのメンバー全てを移動させる(あるいは名札をつけかえる)必要が出てくるからだ。これは効率が良くない。

Union Findにおいては、グループの併合を行う際も、片方のリーダーをもう片方のリーダーの下にくっつけることで、簡単にグループの併合を行うことができる。このやり方は、例えばAをリーダーとするグループとBをリーダーとするグループがある場合、Bに「今からAが上司になるから」と伝えるイメージである。これだけで済むので、非常に効率が良い。

つまり、Union Findとは、友達の友達はみんな友達になるグループ関係において、「xとyが同じグループにいるかどうかを判定する」「xとyを友達にする」という操作を、効率よく行えるようにするデータ構造・探索の仕組みというわけだ。

応用範囲は広く、グラフの連結問題・閉路の発見などはもちろん、例えば「嘘つきと正直者」のようなグループ分け問題に使用することもできる場合がある。

が、いちいち問題が出る度にイチから書くのは面倒なので、あらかじめ用意しておくと吉。というわけで、用意しておくことにしたのが、上記のコードである。
ネット上には他に良いコードがいくらでも落ちているが、冗長性を確保するという意味でも、上げておくことにする。

Share

コメントを残す

メールアドレスが公開されることはありません。