Union_Find (Disjoint Set Union) クラス in Python3

# Union_Findクラス in Python3
# 書いた人: scrblbug
# サイトURL: http://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_member_count(x):xの所属するグループのメンバー数を取得
# .get_group_members(x):xの所属するグループのメンバーをリストで取得
# .get_all_member_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_member_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_member_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):
    # 初期化は要素数0として行うので注意。
    def __init__(self):
        super().__init__(0)
        self.obj_to_idx = {}
        self.idx_to_obj = []
    
    # 新規オブジェクトをUnion Findに独立要素として追加する
    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)

    def obj_existp(self, obj_x):
        return obj_x in self.obj_to_idx