并查集模板

Python版本

class UF:
    parent = {}
    size = {}
    cnt = 0
    def __init__(self, M):
        # 初始化 parent，size 和 cnt
        # self.parent = {i for i in range(n)}

    def find(self, x):
        while x != self.parent[x]:
            x = self.parent[x]
            self.parent[x] = self.parent[self.parent[x]]  # 路径压缩：隔代压缩（基于循环）
        return x
        
    # def find(self, x):
    #     if x != self.parent[x]:
    #         self.parent[x] = self.find(self.parent[x])  # 路径压缩：完全压缩（基于递归）
    #     return self.parent[x]
    
    def union(self, p, q):
        if self.connected(p, q): return
        leader_p = self.find(p)  # 小的树挂到大的树上，使树尽量平衡，即按秩合并，基于size的优化方法
        leader_q = self.find(q)
        if self.size[leader_p] < self.size[leader_q]:
            self.parent[leader_p] = leader_q
        else:
            self.parent[leader_q] = leader_p
        self.cnt -= 1
    def connected(self, p, q):
        return self.find(p) == self.find(q)