如何以优于O(len(set))的复杂度合并分组集合？Union-Find是否适用？

我有一个名为groups的集合列表，groups[i]是与i同组的标签（整数）集合。示例如下：

# idx:     0    1    2        3    4      5        6        7
groups = [{0}, {1}, {2,5,6}, {3}, {4,7}, {2,5,6}, {2,5,6}, {4,7}]

这表示：

• 0属于仅包含0的组
• 1属于仅包含1的组
• 2、5、6属于同一组，即包含2、5、6的组
• 3属于仅包含3的组
• 4和7属于同一组，即包含4、7的组

有时会发现两个组实际是同一组，比如发现4和5同组，就需要合并{2,5,6}和{4,7}，合并后groups应变为：

# idx:     0    1    2            3    4            5            6            7
groups = [{0}, {1}, {2,4,5,6,7}, {3}, {2,4,5,6,7}, {2,4,5,6,7}, {2,4,5,6,7}, {2,4,5,6,7}]

朴素方法需要更新len(groupA)+len(groupB)个元素，效率较低。一种更高效的方式是让groups中存储同一集合的引用：

# idx:     0    1   2        3   4       5       6       7
groups = [{0}, {1}, groupA, {3}, groupB, groupA, groupA, groupB]

其中groupA == {2,5,6}，groupB == {4,7}。合并时最多需要min(len(groupA), len(groupB))次更新：

if len(groupA) < len(groupB):
  groupA, groupB = groupB, groupA
# len(groupA) >= len(groupB)
groupA |= groupB
for label in groupB:
  groups[label] = groupA

合并后groupA变为{2,4,5,6,7}，groups对应位置都指向groupA。

请问是否有更高效的分组合并方式？Union-Find（不相交集合，Disjoint Set Union）数据结构是否适用？

Union-Find（简称DSU）不仅适用，它正是解决这类动态连通性问题的最优选择，比你当前的引用集合方法效率更高。

为什么Union-Find更高效？ #

你当前的方法已经通过引用和合并时只更新小集合元素优化了效率，但Union-Find进一步把合并和查询操作的时间复杂度降到了近乎常数（经过路径压缩和按秩合并优化后，时间复杂度为α(n)，其中α是阿克曼函数的反函数，增长极慢，对于实际应用中的n，α(n)不超过5）。

相比之下，你当前的方法每次合并需要遍历小集合的所有元素（O(min(len(A), len(B)))），而Union-Find的合并操作只需要常数级的父节点更新，查询操作也通过路径压缩快速找到根节点。

Union-Find的核心实现思路 #

Union-Find不需要维护完整的集合内容，只需要维护两个数组：

1. parent数组：parent[i]表示i的父节点，初始时parent[i] = i（每个元素自成一组）。
2. rank（或size）数组：用于记录每个树的高度（或大小），合并时保证小树合并到大树下，避免树结构退化。

核心操作：

1. Find（查找根节点）：
   递归或迭代查找元素的根节点，同时进行路径压缩——把路径上的所有节点直接指向根节点，后续查询会更快。

   def find(u):
       if parent[u] != u:
           parent[u] = find(parent[u])  # 路径压缩
       return parent[u]
   

2. Union（合并两个组）：
   先找到两个元素的根节点，如果根不同，则将秩小的树合并到秩大的树下，更新秩数组。

   def union(u, v):
       root_u = find(u)
       root_v = find(v)
       if root_u == root_v:
           return  # 已经同组，无需合并
       # 按秩合并，把小树合并到大树下
       if rank[root_u] < rank[root_v]:
           parent[root_u] = root_v
       else:
           parent[root_v] = root_u
           if rank[root_u] == rank[root_v]:
               rank[root_u] += 1
   

对比你的当前方法 #

• 你的方法需要维护实际的集合对象，合并时要更新小集合内所有元素的引用；而Union-Find只维护父节点和秩，不需要存储集合内容，内存占用更低。
• 当需要查询两个元素是否同组时，你的方法需要比较引用是否相同，而Union-Find只需要调用find看根节点是否一致，效率更高。
• 对于大规模数据的多次合并/查询操作，Union-Find的优势会更加明显，因为它的时间复杂度几乎是常数级。

如何适配你的场景 #

如果需要获取某个组的所有元素，Union-Find本身不直接存储集合，但可以通过额外维护一个components字典，键是根节点，值是该组的元素列表。合并时把小的列表合并到大的列表中，这样既保留了Union-Find的高效性，又能快速获取组内元素：

components = {i: {i} for i in range(n)}

def union(u, v):
    root_u = find(u)
    root_v = find(v)
    if root_u == root_v:
        return
    if len(components[root_u]) < len(components[root_v]):
        root_u, root_v = root_v, root_u
    parent[root_v] = root_u
    components[root_u].update(components[root_v])
    del components[root_v]

这样合并操作的时间复杂度还是由find的路径压缩决定，而获取组元素可以直接通过根节点查询components字典。

内容的提问来源于stack exchange，提问作者joseville