并查集:用于处理一些不交集(Disjoint sets)的合并及查询问题。
在求最小生成树的 Kruskal 算法 中就用到了并查集。
算法
添加
新添加一个元素时,它的 parent 是它自己。特别地,初始化 union find 时可用
reprs = [i for i in range(n)]
sizes = [1 for i in range(n)]
# 或者
ranks = [1 for i in range(n)]sizes 或 ranks 用来保证 union 的时候把小/矮树接到大/高树上,从而节省时间。
查询
每个集合的代表是集合的 root 节点。在查询时进行路径压缩。
def find(x):
if reprs[x] != x:
reprs[x] = find(reprs[x])
return reprs[x]合并
如果两个节点的代表不同,则将它们的代表相连。注意 merge 时要对两节点的代表(ri 和 rj)进行操作,而不是 i 和 j。
如果使用 sizes,把小树接到大树上,并更新新 root 的 size(两棵树 size 之和):
def merge(i, j):
ri = find(i)
rj = find(j)
if ri != rj:
if sizes[ri] > sizes[rj]:
reprs[rj] = ri
sizes[ri] += sizes[rj]
else:
reprs[ri] = rj
sizes[rj] += sizes[ri]需要注意只有 root 的 size 信息是准确的。查询 size 的时候应该使用 sizes[find(x)]。
如果使用 ranks,只需在 merge 两个 rank 相同的树的时候把新 root 的 rank +1:
def merge(i, j):
ri = find(i)
rj = find(j)
if ri != rj:
if ranks[ri] > ranks[rj]:
reprs[rj] = ri
elif ranks[ri] < ranks[rj]:
reprs[ri] = rj
else:
reprs[ri] = rj
ranks[rj] += 1时间复杂度
使用路径压缩加上 size/rank 可使单次操作平均时间复杂度达到 (反阿克曼函数),接近常数时间。
相关题目
LeetCode 547. Number of Provinces
Python 解法:
def findCircleNum(self, isConnected: List[List[int]]) -> int:
'''
Union Find solution. Begins by seting the number of provinces to the
number of cities, then union each pair of the connected cities. If they
belong to different sets before the union, reduce the number of
provinces by 1 (unioned two cities into one province).
'''
city_num = len(isConnected)
representative = [i for i in range(city_num)]
def find(c):
if representative[c] == c:
return c
else:
representative[c] = find(representative[c])
return representative[c]
ret = city_num
for i in range(city_num):
for j in range(i + 1, city_num):
if not isConnected[i][j]:
continue
if (ri := find(i)) != (rj := find(j)):
# Connected cities are in different sets. Union them and
# reduce the number of provinces
representative[ri] = rj
ret -= 1
return ret