Golang - Optimized Hashmap Tree-based Disjoint Set

tbds.go

package main

import "fmt"

type DisjointSet interface {
	Find(item int) int
	Union(setA int, setB int)
}

type hashmapDisjointSet struct {
	parent map[int]int
	rank   map[int]int // record the depth of the trees representing the sets
}

// Must implement all
var _ DisjointSet = (*hashmapDisjointSet)(nil)

func NewHashmapDisjointSet(universe []int) DisjointSet {
	m := make(map[int]int)
	r := make(map[int]int)

	for _, c := range universe {
		m[c] = c
		r[c] = 0
	}
	// We have len(universe) disjoint sets, each set cointain one item, hence its parent is itself

	// Hardcode some init state, the first set and last set in the universe will be in the same set now
	m[universe[0]] = universe[len(universe)-1]
	r[universe[len(universe)-1]] = 1

	return &hashmapDisjointSet{
		parent: m,
		rank:   r,
	}
}

func (s *hashmapDisjointSet) Find(item int) int {
	if s.parent[item] == item {
		return item
	}

	return s.Find(s.parent[item])
}

func (s *hashmapDisjointSet) Union(setA int, setB int) {
	if s.rank[setA] > s.rank[setB] {
		s.parent[setB] = setA
	} else if s.rank[setB] > s.rank[setA] {
		s.parent[setA] = setB
	} else {
		s.parent[setA] = setB
		s.rank[setB]++
	}
}

func main() {
	universe := []int{1, 2, 3, 4, 5}
	ds := NewHashmapDisjointSet(universe)

	a := ds.Find(2) // returns 2
	ds.Union(2, 3)  // 2 and 3 now in the same set 3
	b := ds.Find(2) // returns 3
	ds.Union(3, 5)  // 1, 2, 3, and 5 now in the same set 5
	c := ds.Find(2) // returns 5

	fmt.Println(a, b, c)
}