带权并查集python实现

以下是带权并查集的 Python 实现：

class WeightedUnionFind:
    def __init__(self, n: int):
        self.parent = list(range(n))
        self.rank = [0] * n
        self.size = [1] * n

    def find(self, x: int) -> int:
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, x: int, y: int, weight: int) -> bool:
        root_x, root_y = self.find(x), self.find(y)
        if root_x == root_y:
            return False
        if self.rank[root_x] < self.rank[root_y]:
            self.parent[root_x] = root_y
            self.size[root_y] += self.size[root_x]
        else:
            self.parent[root_y] = root_x
            self.size[root_x] += self.size[root_y]
            if self.rank[root_x] == self.rank[root_y]:
                self.rank[root_x] += 1
        return True

    def get_size(self, x: int) -> int:
        return self.size[self.find(x)]

    def is_connected(self, x: int, y: int) -> bool:
        return self.find(x) == self.find(y)

其中，`parent` 数组用于记录每个节点的父节点，`rank` 数组用于记录每个节点所在树的深度（如果两棵树合并时深度相同，则将深度加 1），`size` 数组用于记录每个集合的元素个数。

`find` 函数用于查找某个节点所属的集合，如果该节点不是根节点，则将其父节点设置为根节点并返回根节点。

`union` 函数用于合并两个集合，并返回是否成功合并。首先查找两个节点所在树的根节点，如果它们已经在同一个集合中，则返回 False。否则，将深度较小的树合并到深度较大的树下，并更新 `size` 数组、`rank` 数组等信息。

`get_size` 函数用于获取某个节点所在集合的元素个数，直接返回该节点所在集合的 `size` 值即可。

`is_connected` 函数用于判断两个节点是否在同一个集合中，直接比较它们的根节点是否相同即可。

使用示例：

uf = WeightedUnionFind(5)
uf.union(0, 1, 2)
uf.union(1, 2, 3)
uf.union(3, 4, 1)
print(uf.is_connected(0, 4))   # True
print(uf.get_size(1))           # 4

输出：

True
4