并查集（UnionFind）技巧总结

什么是并查集

在计算机科学中，并查集是一种树型的数据结构，用于处理一些不交集（Disjoint Sets）的合并及查询问题。有一个联合-查找算法（Union-find Algorithm）定义了两个用于此数据结构的操作：

• Find：确定元素属于哪一个子集。它可以被用来确定两个元素是否属于同一子集。
• Union：将两个子集合并成同一个集合。

由于支持这两种操作，一个不相交集也常被称为联合-查找数据结构（Union-find Data Structure）或合并-查找集合（Merge-find Set）。

并查集可以解决什么问题

• 组团、配对
• 图的连通性问题
• 集合的个数
• 集合中元素的个数

算法模板

type UnionFind struct {
	count  int
	parent []int
}

func ctor(n int) UnionFind {
	uf := UnionFind{
		count:  n,
		parent: make([]int, n),
	}
	for i := 0; i < n; i++ {
		uf.parent[i] = i
	}
	return uf
}

func (uf *UnionFind) find(p int) int {
	for p != uf.parent[p] {
		uf.parent[p] = uf.parent[uf.parent[p]] // 路径压缩
		p = uf.parent[p]
	}
	return p
}

func (uf *UnionFind) union(p, q int) {
	rootP, rootQ := uf.find(p), uf.find(q)
	if rootP == rootQ {
		return
	}
	uf.parent[rootP] = rootQ
	uf.count--
}

实例

547. 朋友圈

题目分析：

题目求的是有多少个朋友圈，也就是求有集合个数，可用并查集解决。

两重遍历所有学生，判断俩俩是否为朋友，如为朋友将加入到集合中。这里可以通过遍历二维矩阵的右半边即可，可降低遍历数量，从而降低时间复杂度。

代码实现：

func findCircleNum(M [][]int) int {
	n := len(M)
	uf := ctor(n)
	// 遍历学生 i, j ，if M[i][j]==1 加入集
	for i := 0; i < n; i++ {
		for j := i + 1; j < n; j++ {
			if M[i][j] == 1 {
				uf.union(i, j)
			}
		}
	}
	// 再返回有多少个集合
	return uf.count
}

type UnionFind struct {
	parents []int
	count   int
}

func ctor(n int) UnionFind {
	uf := UnionFind{
		parents: make([]int, n),
		count:   n,
	}

	for i := 0; i < n; i++ {
		uf.parents[i] = i
	}

	return uf
}

func (uf *UnionFind) find(p int) int {
	for p != uf.parents[p] {
		uf.parents[p] = uf.parents[uf.parents[p]]
		p = uf.parents[p]
	}
	return p
}

func (uf *UnionFind) union(p, q int) bool {
	rootP, rootQ := uf.find(p), uf.find(q)
	if rootP == rootQ {
		return false
	}
	uf.parents[rootP] = rootQ
	uf.count--
	return true
}

复杂度分析：

• 时间复杂度：(O(n^2))。两重遍历用时 (O(n^2)) ，uf.union 和 uf.find 的时间复杂度为 (O(1)) ，所以总的时间复杂度为 (O(n^2))。
• 空间复杂度：(O(n))。需要一个 (O(n)) 大小的空间。

200. 岛屿数量

题目分析：

题目求的是岛屿数量，即集合个数，可用并查集解决。

题目可抽象为遍历所有网格 (i, j)，如果是陆地（(i, j) == '1'），则把其右边的陆地（(i+1, j) == '1'）和下边的陆地（(i, j+1) == '1'）合并到一起；如是水（(i, j) == '0'），则把其合并到一个哨兵集合里。最后返回 集合个数 - 1。

注：这里关键是对于水的处理，把其合并到一个哨兵集合里，让水不会单独存在，从而干扰岛屿个数的判断。

代码实现：

func numIslands(grid [][]byte) int {
	rows := len(grid)
	if rows == 0 {
		return 0
	}
	cols := len(grid[0])
	if cols == 0 {
		return 0
	}

	uf := ctor(rows*cols + 1)
	guard := rows * cols // 哨兵：用于作为 '0' 的集合
	directions := [][]int{[]int{0, 1}, []int{1, 0}}

	for i := 0; i < rows; i++ {
		for j := 0; j < cols; j++ {
			index := i*cols + j
			if grid[i][j] == '1' {
				for _, direction := range directions {
					newI, newJ := i+direction[0], j+direction[1]
					if newI < rows && newJ < cols && grid[newI][newJ] == '1' {
						newIndex := newI*cols + newJ
						uf.union(index, newIndex)
					}
				}
			} else {
				uf.union(guard, index)
			}
		}
	}
	return uf.count - 1
}

type UnionFind struct {
	parents []int
	count   int
}

func ctor(n int) UnionFind {
	uf := UnionFind{parents: make([]int, n), count: n}
	for i := 0; i < n; i++ {
		uf.parents[i] = i
	}
	return uf
}

func (uf *UnionFind) find(p int) int {
	for p != uf.parents[p] {
		uf.parents[p] = uf.parents[uf.parents[p]]
		p = uf.parents[p]
	}
	return p
}

func (uf *UnionFind) union(p, q int) bool {
	rootP, rootQ := uf.find(p), uf.find(q)
	if rootP == rootQ {
		return false
	}
	uf.parents[rootP] = rootQ
	uf.count--
	return true
}

复杂度分析：

• 时间复杂度：(O(n*m))，其中 n 和 m 分别表示二维数组的行数和列数。
• 空间复杂度：(O(n*m))。并查集需要 n * m 大小的数组空间。

130. 被围绕的区域

题目分析：

题目可理解为把边界上的 'O' 保留，其他都填充为 'X' ，可以把边界上的 'O' 作为一个集合，不是这个集合的填充为 'X' ，因此可使用并查集解决。

1. 遍历边界上的点，把 'O' 合并到一个哨兵集合里。
2. 遍历二维矩阵里的点，把 'O' 的右和下合并到一起。
3. 遍历二维矩阵，把不在哨兵集合里的全部填充为 'X'

代码实现：

func solve(board [][]byte) {
	n := len(board)
	if n == 0 {
		return
	}
	m := len(board[0])
	if m == 0 {
		return
	}
	uf := ctor(n*m + 1)
	guard := n * m
	directions := [][]int{[]int{0, 1}, []int{1, 0}}

	getIndex := func(i, j int) int {
		return i*m + j
	}
	// 1. 遍历边界上的点，把 'O' 合并到一个哨兵集合里。
	for j := 0; j < m; j++ {
		if board[0][j] == 'O' {
			uf.union(getIndex(0, j), guard)
		}
		if board[n-1][j] == 'O' {
			uf.union(getIndex(n-1, j), guard)
		}
	}
	for i := 0; i < n; i++ {
		if board[i][0] == 'O' {
			uf.union(getIndex(i, 0), guard)
		}
		if board[i][m-1] == 'O' {
			uf.union(getIndex(i, m-1), guard)
		}
	}

	// 2. 遍历二维矩阵里的点，把 ```'O'``` 的右和下合并到一起。
	for i := 0; i < n; i++ {
		for j := 0; j < m; j++ {
			if board[i][j] == 'O' {
				for _, direction := range directions {
					newI, newJ := i+direction[0], j+direction[1]
					if newI < n && newJ < m && board[newI][newJ] == 'O' {
						uf.union(getIndex(newI, newJ), getIndex(i, j))
					}
				}
			}
		}
	}

	// 3. 遍历二维矩阵，把不在哨兵集合里的全部填充为 'X'
	for i := 0; i < n; i++ {
		for j := 0; j < m; j++ {
			if !uf.isConnect(getIndex(i, j), guard) {
				board[i][j] = 'X'
			}
		}
	}
}

type UnionFind struct {
	parents []int
	count   int
}

func ctor(n int) UnionFind {
	uf := UnionFind{
		parents: make([]int, n),
		count:   n,
	}

	for i := 0; i < n; i++ {
		uf.parents[i] = i
	}

	return uf
}

func (uf *UnionFind) find(p int) int {
	for p != uf.parents[p] {
		uf.parents[p] = uf.parents[uf.parents[p]]
		p = uf.parents[p]
	}
	return p
}

func (uf *UnionFind) union(p, q int) bool {
	rootP, rootQ := uf.find(p), uf.find(q)
	if rootP == rootQ {
		return false
	}
	uf.parents[rootP] = rootQ
	uf.count--
	return true
}

func (uf *UnionFind) isConnect(p, q int) bool {
	return uf.find(p) == uf.find(q)
}

复杂度分析：

• 时间复杂度：(O(n^2))，其中 n 和 m 分别表示二维数组的行数和列数。
• 空间复杂度：(O(n^2))。并查集需要 n * m 大小的数组空间。

总结

1. 要熟练掌握并查集的模板，要能够快速写出来。
2. 要掌握并查集的应用场景。例如组团、配对、图的连通性问题、集合个数、集合中元素的个数等。
3. 对于二维的问题转一维解决，例如 200. 岛屿数量 和 130. 被围绕的区域。
4. 找出元素间的“配对”关系是解决问题的关键。例如二维数组，找当前位置与其右和其下配对。例如 200. 岛屿数量 和 130. 被围绕的区域。

参考资料

• 并查集 - 力扣
• 并查集概念及用法分析