github.com/qiuhoude/go-web@v0.0.0-20220223060959-ab545e78f20d/algorithm/datastructures/unionfind/unionfind.go (about) 1 // 并查集 2 package unionfind 3 4 import "errors" 5 6 type Uf interface { 7 GetSize() int 8 IsConnected(p, q int) bool // p,q两个下标是是否连接 9 UnionElements(p, q int) error // p,q下标进行合并 10 } 11 12 type UnionFind struct { 13 parent []int // parent[i]表示第i个元素所指向的父节点 14 rank []int //rank[i]表示以i为根的集合所表示的树的层数,在路径压缩之后不表示层数而是一个分数 15 } 16 17 var ArgumentErr = errors.New("index is out of bound") 18 19 func NewUnionFind(size int) *UnionFind { 20 p := make([]int, size) 21 rank := make([]int, size) 22 for i := 0; i < size; i++ { 23 p[i] = i 24 rank[i] = 1 25 } 26 return &UnionFind{parent: p} 27 } 28 29 func (u *UnionFind) GetSize() int { 30 return len(u.parent) 31 } 32 33 // 查找过程, 查找元素p所对应的集合编号 34 // O(h)复杂度, h为树的高度 35 func (u *UnionFind) find(p int) (int, error) { 36 if p < 0 || p > u.GetSize() { 37 return 0, ArgumentErr 38 } 39 // 不断去寻找自己的父节点, 直到到达根节点 40 // 根节点的特点: parent[p] == p 41 for p != u.parent[p] { 42 u.parent[p] = u.parent[u.parent[p]] // 将自己的父节点 指向爷节点 43 p = u.parent[p] // 直接从爷节点开始往上 44 } 45 46 return p, nil 47 } 48 49 func (u *UnionFind) UnionElements(p, q int) error { 50 pRoot, err := u.find(p) 51 if err != nil { 52 return err 53 } 54 qRoot, err := u.find(p) 55 if err != nil { 56 return err 57 } 58 if qRoot == pRoot { 59 return nil 60 } 61 62 // 根据两个元素所在树的rank不同判断合并方向 63 // 将rank低的集合合并到rank高的集合上 64 if u.rank[pRoot] < u.rank[qRoot] { 65 u.parent[pRoot] = qRoot 66 } else if u.rank[pRoot] > u.rank[qRoot] { 67 u.parent[qRoot] = pRoot 68 } else { // rank 相等 69 u.parent[pRoot] = qRoot 70 u.rank[qRoot] += 1 71 } 72 u.parent[pRoot] = qRoot 73 return nil 74 } 75 76 func (u *UnionFind) IsConnected(p, q int) bool { 77 pRoot, err := u.find(p) 78 if err != nil { 79 return false 80 } 81 qRoot, err := u.find(p) 82 if err != nil { 83 return false 84 } 85 return pRoot == qRoot 86 }