算法基础-并查集

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40

func NewUF(n int) *UnionFind {
	parent := make([]int, n)
	for i := range parent {
		parent[i] = i
	}
	return &UnionFind{
		parent: parent,
		count:  n,
	}
}

type UnionFind struct {
	parent []int
	count  int
}

func (uf *UnionFind) Union(a, b int) {
	rootA := uf.Find(a)
	rootB := uf.Find(b)
	if rootA == rootB {
		return
	}
	uf.parent[rootA] = rootB
	uf.count--
}

func (uf *UnionFind) Find(a int) int {
	for uf.parent[a] != a {
		// compress parent tree
		uf.parent[a] = uf.parent[uf.parent[a]]
		a = uf.parent[a]
	}
	return a
}

func (uf *UnionFind) Connected(a, b int) bool {
	rootA := uf.Find(a)
	rootB := uf.Find(b)
	return rootA == rootB
}