pythonでUnionFind（強化版: 要素に文字列やtuple可）

はじめに

かの有名なUnionFindですが、集合の要素に文字列やtupleを扱えるものが少なかったので記事にしてみようと思います。役立つ場面にはあまり遭遇していないので、役立つかどうかはわかりませんが、興味がある方は見ていってください。

参考記事

https://note.nkmk.me/python-union-find/
https://qiita.com/white1107/items/52fd4149bb1846862e38

コードの紹介

UnionFindがどういうものかについては他の記事に任せ、さっそくコードを紹介しようと思います。

UnionFind.py

import collections
class UnionFind():
    def __init__(self):
        '''
        unionfind経路圧縮あり,要素にtupleや文字列可,始めに要素数指定なし
        '''
        self.parents = dict()                                      # {子要素:親ID,}
        self.members_set = collections.defaultdict(lambda : set()) # keyが根でvalueが根に属する要素要素(tupleや文字列可)
        self.roots_set = set()                                     # 根の集合(tupleや文字列可)
        self.key_ID = dict()                                       # 各要素にIDを割り振る
        self.ID_key = dict()                                       # IDから要素名を復元する
        self.cnt = 0                                               # IDのカウンター

    def dictf(self,x): # 要素名とIDをやり取りするところ
        if x in self.key_ID:
            return self.key_ID[x]
        else:
            self.cnt += 1
            self.key_ID[x] = self.cnt
            self.parents[x] = self.cnt
            self.ID_key[self.cnt] = x
            self.members_set[x].add(x)
            self.roots_set.add(x)
            return self.key_ID[x]

    def find(self, x):
        ID_x = self.dictf(x)
        if self.parents[x] == ID_x:
            return x
        else:
            self.parents[x] = self.key_ID[self.find(self.ID_key[self.parents[x]])]
            return self.ID_key[self.parents[x]]

    def union(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if self.parents[x] > self.parents[y]:
            x, y = y, x
        if x == y:
            return
        for i in self.members_set[y]:
            self.members_set[x].add(i)
        self.members_set[y] = set()
        self.roots_set.remove(y)
        self.parents[y] = self.key_ID[x]

    def size(self, x):#xが含まれる集合の要素数
        return len(self.members_set[self.find(x)])

    def same(self, x, y):#同じ集合に属するかの判定
        return self.find(x) == self.find(y)

    def members(self, x):#xを含む集合の要素
        return self.members_set[self.find(x)]

    def roots(self):#根の要素
        return self.roots_set

    def group_count(self):#根の数
        return len(self.roots_set)

    def all_group_members(self):#根とその要素
        return {r: self.members_set[r] for r in self.roots_set}

if __name__=="__main__": #デバック用
    uf_s = UnionFind()
    uf_s.union('A', 'D')
    uf_s.union('D', 'C')
    uf_s.union('E', 'B')
    print(uf_s.members('D'))
    print(uf_s.members('E'))
    uf_s.union('A', 'B')
    print(uf_s.members('E'))
    print(uf_s.same('A','F'))
    print(uf_s.members('F'))
    print(uf_s.members('G'))

    print()
    uf_t = UnionFind()
    uf_t.union((2,3),(3,4))
    uf_t.union((4,3),(5,4))
    uf_t.union((1,3),(3,4))
    print(uf_t.members((2,3)))
    print(uf_t.roots())
    print(uf_t.group_count())
    print(uf_t.all_group_members())

    print()
    uf_ts = UnionFind()
    uf_ts.union((2,3),1)
    uf_ts.union((4,3),'A')
    uf_ts.union((1,3),(3,4))
    uf_ts.union((1,3),'A')
    print(uf_ts.members((2,3)))
    print(uf_ts.roots())
    print(uf_ts.all_group_members())
    print(uf_ts.size('A'))

テキトーなデバック用コードを下につけておきました。

説明

使い方

基本的にはよく見るUnionFind同じ使い方ができます。
find(x)：xの根の取得
union(x,y)：集合を合体
size(x)：x要素が含まれる集合の大きさを出力
same(x,y)：同じ集合に属するか判定
members(x)：xが含まれる集合を出力
roots()：すべての根を取得
group_count()：集合の数を出力
all_group_members()：すべての根とその要素を取得

工夫点と利点

とにかく辞書を駆使して、要素と要素番号の対応を取りながら経路圧縮も行い集合を管理しています。なので辞書のkeyにできる要素なら型が違っていても同じように扱うことができます（文字列、tuple、整数などなど）。また、他で見るunionfindは要素数を先に指定するものが多いですが、本コードでは要素がやってくるたび初めてかどうかを確認し、初めてなら今まで使っていない番号を割り当てるようにしているので要素数の指定は不要です。
(追記:座標圧縮が必要な時や二次元unionfindやりたいときにも使えますね。)

競プロ典型 90 問　012 - Red Painting（★4）
二次元UnionFindを使いたい問題ですが、このUnionfindにそのまま突っ込めばACでした(提出結果)。

弱点

普通のunionfindに比べ、辞書等で管理している部分が多いのでメモリの消費が大きいです。また、unionfind内部でのデータのやり取りが多く計算時間も多くかかってしまいます。
例：https://atcoder.jp/contests/practice2/tasks/practice2_a
TLEはしませんでしたが、1500msくらいでこれ以上重くなると通らないことがありそうです。

おわりに

ここまでお読みいただきありがとうございました。私自身競技プログラミング初学者ですので、間違い等ございましたらご指摘よろしくお願いいたします。（加えて、本コードが役立つ問題がありましたらリンク等コメントください。）