Python Union-Find 并查集

Question

并查集，顾名思义，就是实现集合的并、以及查元素属于哪个集合的功能。比如以下并查集：

{0} {1,5,8} {2,3,4,6} {7,9}

常用操作有：

• 查找元素属于哪个集合 Find
• 合并两个集合 Union：比如 Union(1, 7) 就是将 {1,5,8} 和 {7,9} 两个集合合并为 {1,5,8,7,9}
• 判断两个元素是否在同一个集合 Connected

实现方式1 -- 线性结构

一种最简单的实现并查集的方式是使用线性结构，用数组存储，数组的每个位置存储元素的值，以及该元素所属的集合ID，这个ID是集合中某个元素的数组下标。这样实现Find的时候直接返回该元素的集合ID，实现Union时先判断两个集合ID是否相同，若不同，则将其中一个集合的ID全部换为另一个集合的ID；实现Connected只需要判断两个元素的集合ID是否相同。

但是这种实现方法的一个问题就是Union的时间复杂度过高，Union需要把一个集合中的每一个元素的ID都替换掉，当元素个数很多时，替换次数就会很多

实现方式2 -- 树形结构

使用树形结构（不是树这种数据结构）表示并查集，每个结点代表一个集合元素，含有一个指向父结点的指针。使用数组实现，数组中每个位置存放该元素的值和父结点在数组中的索引（根结点存放的索引就是自身的索引）。

在树形结构实现的并查集中，Find的实现依然是遍历数组，返回元素的根结点；Union只需要找到两个元素的根结点，如果不同根，将其中一个元素的根结点指向的父结点设置成另一个元素根结点的数组索引。因此Union的复杂度变为了树的高度O(lg n)

树形结构的改进 -- 让树更加平衡

使用树形结构解决了Union操作复杂度高的问题，但是还可以继续改进。可以发现当合并两个集合时，是直接将一颗树接到另一个树的根结点上，这样有可能会导致某棵树不平衡，最坏情况下时间复杂度变为线性。解决方法就是在执行Union操作之前，判断要合并的两棵树，只将小的树合并到大的树中，这样可以保持一定的平衡性，将最坏情况下的时间复杂度降至O(lg n)。这种方法需要一个新的数组存储每棵树的大小。

下面用代码实现这种并查集：

class UnionFind:
  # 初始化并查集中的元素，默认每个元素的值就是数组下标，每个元素的集合只有自身
  def __init__(self, N):
    self.data = list(range(N))  # 存储元素的值
    self.parent = list(range(N))  # 存储该元素的父结点的数组下标
    self.size = [1] * N  # 存储以该元素为根的树的大小（树中的元素个数）
    self.count = N  # 存储集合的个数

find(value) -- 返回元素的根结点的数组下标

def find(self, value):  # 返回根结点的下标
  index = self.data.index(value)
  while self.parent[index] != index:
    index = self.parent[index]
  return index

union(value1, value2) -- 合并两个元素所属的集合

def union(self, value1, value2):
  root1 = self.find(value1)
  root2 = self.find(value2)
  if root1 == root2:
    return
  if self.size[root1] > self.size[root2]:
    self.parent[root2] = root1
    self.size[root1] += self.size[root2]
  else:
    self.parent[root1] = root2
    self.size[root2] + self.size[root1]
  self.count -= 1

connected(value1, value2) -- 判断两个元素是否属于同一集合

def connected(self, value1, value2):
  return self.find(value1) == self.find(value2)