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  }