功能说明
并查集用于维护一组不相交集合,支持“合并两个元素所在集合”和“判断两个元素是否连通”。这段代码实现了路径压缩与按秩合并两个优化,适合处理账号去重、网络连通性、图片连通区域、Kruskal 最小生成树等场景。
components 字段实时记录当前连通分量数量,每次成功合并减一;如果两个元素已经在同一集合,union 返回 False,不会错误地减少计数。
关键代码
find 先用循环找到根节点,再第二次遍历把路径上的每个节点直接挂到根上,这就是路径压缩。它使用迭代而不是递归,避免深链导致递归深度超限。
union 先找到两个根;相同则返回 False。随后按 rank 把较矮的树挂到较高的树下,只有秩相等时才增加新根的秩,从而把树高控制在近似对数级别。connected 只是比较两个根是否相同。
使用注意
并查集只支持合并,不支持拆分;如果需要删除元素或撤销合并,应使用可持久化并查集或重建结构。所有节点必须在初始化时登记,访问未登记的键会触发 KeyError。
该实现不是线程安全的,多线程合并需要外部加锁或使用并发数据结构。秩是启发式上界,不表示精确高度。元素类型假定为字符串,实际可泛化为任意可哈希对象。
验证方式
示例依次合并 a-b、b-c、d-e,前三次都返回 True;a 与 c 连通、a 与 d 不连通;分量数为 2。再次合并 a-c 返回 False,因为二者已经连通。可以再加入单元素集合,确认初始分量为 1 且查询自身返回 True。
代码示例
from __future__ import annotations
class DisjointSet:
def __init__(self, items: list[str]) -> None:
self.parent = {item: item for item in items}
self.rank = {item: 0 for item in items}
self.components = len(self.parent)
def find(self, item: str) -> str:
root = item
while self.parent[root] != root:
root = self.parent[root]
while self.parent[item] != item:
parent = self.parent[item]
self.parent[item] = root
item = parent
return root
def union(self, left: str, right: str) -> bool:
root_left = self.find(left)
root_right = self.find(right)
if root_left == root_right:
return False
if self.rank[root_left] < self.rank[root_right]:
root_left, root_right = root_right, root_left
self.parent[root_right] = root_left
if self.rank[root_left] == self.rank[root_right]:
self.rank[root_left] += 1
self.components -= 1
return True
def connected(self, left: str, right: str) -> bool:
return self.find(left) == self.find(right)
def main() -> None:
network = DisjointSet(["a", "b", "c", "d", "e"])
for left, right in [("a", "b"), ("b", "c"), ("d", "e")]:
print(left, right, network.union(left, right))
print(network.connected("a", "c"), network.connected("a", "d"))
assert network.components == 2
assert not network.union("a", "c")
if __name__ == "__main__":
main()