59 lines
1,020 B
Go
59 lines
1,020 B
Go
package unionfind
|
|
|
|
type weightedQuickUnion struct {
|
|
id []int
|
|
sizes []int
|
|
count int
|
|
}
|
|
|
|
func NewWeightedQuickUnion(count int) UnionFind {
|
|
id := make([]int, count)
|
|
sizes := make([]int, count)
|
|
|
|
for i := range id {
|
|
id[i] = i
|
|
}
|
|
|
|
for i := range sizes {
|
|
sizes[i] = 1
|
|
}
|
|
|
|
return &weightedQuickUnion{id, sizes, count}
|
|
}
|
|
|
|
func (qf *weightedQuickUnion) Find(site int) int {
|
|
for site != qf.id[site] {
|
|
site = qf.id[site]
|
|
}
|
|
|
|
return site
|
|
}
|
|
|
|
func (qf *weightedQuickUnion) Union(aSite, bSite int) {
|
|
aComponent := qf.Find(aSite)
|
|
bComponent := qf.Find(bSite)
|
|
|
|
if aComponent == bComponent {
|
|
return
|
|
}
|
|
|
|
aSize := qf.sizes[aComponent]
|
|
bSize := qf.sizes[bComponent]
|
|
|
|
if aSize > bSize {
|
|
qf.id[bComponent] = aComponent
|
|
qf.sizes[aComponent] += bSize
|
|
} else {
|
|
qf.id[aComponent] = bComponent
|
|
qf.sizes[bComponent] += aSize
|
|
}
|
|
qf.count--
|
|
}
|
|
|
|
func (qf *weightedQuickUnion) Connected(aSite, bSite int) bool {
|
|
return qf.Find(aSite) == qf.Find(bSite)
|
|
}
|
|
|
|
func (qf *weightedQuickUnion) Count() int {
|
|
return qf.count
|
|
}
|