github.com/qiuhoude/go-web@v0.0.0-20220223060959-ab545e78f20d/algorithm/datastructures/unionfind/optimizesize.go (about)

     1  package unionfind
     2  
     3  // 基于size的优化,将集
     4  type UnionFindOpSize struct {
     5  	parent []int
     6  	sz     []int
     7  }
     8  
     9  func NewUnionFindOpSize(size int) *UnionFindOpSize {
    10  	p := make([]int, size)
    11  	sz := make([]int, size)
    12  	for i := 0; i < size; i++ {
    13  		p[i] = i
    14  		sz[i] = 1
    15  	}
    16  	ret := &UnionFindOpSize{sz: sz}
    17  	ret.parent = p
    18  	return ret
    19  }
    20  
    21  func (u *UnionFindOpSize) GetSize() int {
    22  	return len(u.parent)
    23  }
    24  
    25  // 查找过程, 查找元素p所对应的集合编号
    26  // O(h)复杂度, h为树的高度
    27  func (u *UnionFindOpSize) find(p int) (int, error) {
    28  	if p < 0 || p > u.GetSize() {
    29  		return 0, ArgumentErr
    30  	}
    31  	// 不断去寻找自己的父节点, 直到到达根节点
    32  	// 根节点的特点: parent[p] == p
    33  	for p != u.parent[p] {
    34  		p = u.parent[p]
    35  	}
    36  
    37  	return p, nil
    38  
    39  	// 递归算法
    40  	//if p != u.parent[p] {
    41  	//	u.parent[p], _ = u.find(u.parent[p]) // 将找到的父节点的父节点复制给父节点
    42  	//}
    43  	//return u.parent[p], nil
    44  }
    45  
    46  func (u *UnionFindOpSize) IsConnected(p, q int) bool {
    47  	pRoot, err := u.find(p)
    48  	if err != nil {
    49  		return false
    50  	}
    51  	qRoot, err := u.find(p)
    52  	if err != nil {
    53  		return false
    54  	}
    55  	return pRoot == qRoot
    56  }
    57  
    58  func (u *UnionFindOpSize) UnionElements(p, q int) error {
    59  	pRoot, err := u.find(p)
    60  	if err != nil {
    61  		return err
    62  	}
    63  	qRoot, err := u.find(p)
    64  	if err != nil {
    65  		return err
    66  	}
    67  	if qRoot == pRoot {
    68  		return nil
    69  	}
    70  	// 根据两个元素所在树的元素个数不同判断合并方向
    71  	// 将元素个数少的集合合并到元素个数多的集合上
    72  	if u.sz[pRoot] < u.sz[qRoot] {
    73  		u.parent[pRoot] = qRoot
    74  		u.sz[qRoot] += u.sz[pRoot]
    75  	} else {
    76  		u.parent[qRoot] = pRoot
    77  		u.sz[pRoot] += u.sz[qRoot]
    78  	}
    79  	return nil
    80  }