547.number.of.provinces.disjoint.set.with.path.compression

main.go

package main

import "fmt"

type DisjointSet struct {
	Nodes       []int
	RootsNumber *int
}

func (ds DisjointSet) Find(n int) int {
	if ds.Nodes[n] == n {
		return n
	}
	ds.Nodes[n] = ds.Find(ds.Nodes[n])
	return ds.Nodes[n]
}

func (ds DisjointSet) Union(n1, n2 int) {
	r1 := ds.Find(n1)
	r2 := ds.Find(n2)
	if r1 != r2 {
		for i, v := range ds.Nodes {
			if v == r1 {
				ds.Nodes[i] = r2
			}
		}
		*ds.RootsNumber -= 1
	}
}

func NewDisjointSet(size int) *DisjointSet {
	ds := &DisjointSet{}
	ds.Nodes = make([]int, size)
	for i := 0; i < size; i++ {
		ds.Nodes[i] = i
	}
	ds.RootsNumber = &size
	return ds
}

func main() {
	c := [][]int{{1, 1, 1}, {1, 1, 1}, {1, 1, 1}}
	ds := NewDisjointSet(len(c))
	for i := 0; i < len(c); i++ {
		for j := 0; j < len(c[i]); j++ {
			if c[i][j] == 1 {
				ds.Union(i, j)
			}
		}
	}
	fmt.Printf("len=%d cap=%d nodes=%v size=%v\n", len(ds.Nodes), cap(ds.Nodes), ds.Nodes, *ds.RootsNumber)
}