UnionFind木のpythonでの実装

UnionFind木で同じグループのものを判定する

class UnionFind(object):
    def __init__(self, size):
        self.parent = [i for i in range(size)]
        self.rank = [0 for _ in range(size)]


    def find(self, x):
        if self.parent[x] == x:
            return x
        else:
            return self.find(self.parent[x])


    def unite(self, x, y):
        x = self.find(x)
        y = self.find(y)
        if x == y:
            return

        if self.rank[x] < self.rank[y]:
            self.parent[x] = y
        else:
            self.parent[y] = x
            if self.rank[x] == self.rank[y]:
                self.rank[x] += 1

    def same(self, x, y):
        return self.find(x) == self.find(y)

In [113]: u = UnionFind(10)

In [114]: u.unite(1,2)

In [115]: u.unite(3,2)

In [116]: u.unite(5,6)

In [117]: u.parent
Out[117]: [0, 1, 1, 1, 4, 5, 5, 7, 8, 9]

In [119]: u.rank
Out[119]: [0, 1, 0, 0, 0, 1, 0, 0, 0, 0]