【UnionFind】クラスの実装とAtCoder典型問題まとめ【Python】

# No0
from collections import defaultdict

class UnionFind():
    """
    Union Find木クラス

    Attributes
    --------------------
    n : int
        要素数
    patents : list
        指定した要素の親（1つ上の）要素を格納
        指定した要素が根の場合は，
        -(グループの要素数)  を格納
        => sizeメソッドに反映
    """
    def __init__(self, n):
        self.n = n
        self.parents = [-1] * n

    def find(self, x):
        """
        ノードxの根を見つける
        """
        if self.parents[x] < 0:
            return x
        else:
            self.parents[x] = self.find(self.parents[x])
            return self.parents[x]

    def union(self, x, y):
        """
        木に新たな要素を併合（マージ）

        Parameters
        ---------------------
        x, y : int
            併合するノード
        """
        x = self.find(x)
        y = self.find(y)

        if x == y:
            return

        if self.parents[x] > self.parents[y]:
            x, y = y, x

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

    def size(self, x):
        """
        xの属する木のサイズ
        """
        return -self.parents[self.find(x)]

    def same(self, x, y):
        """
        x, yが同じ木に属するか判定
        """
        return self.find(x) == self.find(y)

    def members(self, x):
        """
        xの属する木に属する要素をリストで返す
        """        
        root = self.find(x)
        return [i for i in range(self.n) if self.find(i) == root]

    def roots(self):
        """
        全ての根をリストで返す
        """ 
        return [i for i, x in enumerate(self.parents) if x < 0]

    def group_count(self):
        """
        グループの数を返す
        """ 
        return len(self.roots())

    def all_group_members(self):
        """
        全てのグループの要素情報を辞書で返す
        """         
        group_members = defaultdict(list)
        for member in range(self.n):
            group_members[self.find(member)].append(member)
        return group_members

    def __str__(self):
        """
        print(uf)で全てのグループの要素情報を簡単に出力する
        """ 
        return '\n'.join(f'{r}: {m}' for r, m in self.all_group_members().items())

n, q = map(int, input().split())
uf = UnionFind(n)
for _ in range(q):
    p, a, b = map(int, input().split())
    if p == 0:
        uf.union(a, b)
    else:
        print("Yes" if uf.same(a, b) else "No")

n, m = map(int, input().split())
 
uf = UnionFind(n)
for i in range(m):
    a, b = map(int, input().split())
    a -= 1
    b -= 1
    uf.union(a, b)
    
print(max([uf.size(i) for i in uf.roots()]))

8
1 5 3 2 5 2 3 1

11
5 1 5 3 2 5 2 3 1 7 2

n = int(input())
a = list(map(int, input().split()))

es = set()
uf = UnionFind(2 * 10 ** 5) ###※(n)だと9ケースでエラーが発生
for i in range(n // 2):
    if a[i] != a[n - 1 - i]:
        uf.union(a[i] - 1, a[n - 1 - i] - 1)

ans = 0
for r in uf.roots():
    if uf.size(r) >= 2:
        ans += uf.size(r) - 1

print(ans)

n, m = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
uf = UnionFind(n)
for i in range(m):
    c, d = map(int, input().split())
    c -= 1
    d -= 1
    uf.union(c, d)
    
flag = True
for c in uf.all_group_members().values():
    if sum(a[i] for i in c) != sum(b[i] for i in c):
        flag = False
        break
print("Yes" if flag else "No")

from collections import defaultdict
n, k, l = map(int, input().split())

uf1 = UnionFind(n)
for i in range(k):
    p, q = map(int, input().split())
    uf1.union(p - 1, q - 1)
    
uf2 = UnionFind(n)
for i in range(l):
    r, s = map(int, input().split())
    uf2.union(r - 1, s - 1)
    
pairs = defaultdict(int)
# 各都市iについて，道路網および鉄道網において
# どの根のグループに属するかを調べる
for i in range(n):
    pairs[(uf1.find(i), uf2.find(i))] += 1
    
for i in range(n):
    print(pairs[(uf1.find(i), uf2.find(i))], end=" ")

n, m = map(int, input().split())

# 辺（エッジの全データを格納）
ed = []
# インスタンスを生成
uf = UnionFind(n)
for _ in range(m):
    # 0-indexed
    x, y = map(int, input().split())
    x, y = x - 1, y - 1
    # 新たなノードとエッジを追加
    uf.union(x, y)
    ed.append((x, y))
    
# 根をrとする連結成分（グループ）の辺の数を格納
es = {r: 0 for r in uf.roots()}
for x, y in ed:
    # ノードx（yでも良い）の根を見つける
    r = uf.find(x)
    # 根rのグループに辺の数を+1
    es[r] += 1

cond = all(uf.size(l) == es[l] for l in uf.roots())
print("Yes" if cond else "No")