# Disjoint sets with path compression to keep track of component number

I wanted to write an implementation of Disjoint Sets with path compression and ranking based on the components' size. In my general approach I followed the description of the data structure in a lecture note I found.

I will use this structure later in another component for connected components, where I want to merge connected cells in a matrix and keep track of the minimum label for each component.

Therefore I added another map that keeps track of the minimum label for each component. An example of what I want to do - consider a matrix of 4-connected components like this (note position 0/0 is the top left corner):

$$\left[ \begin{array}{ccc} 1&2&2\\ 1&3&2\\ 3&3&4\\ \end{array} \right]$$

Now we make position 0/0 (has the label 1) the parent of position 0/1 (has the label 2). Now a) both positions are in the same component and b) both are tracked with label equal to 1, because 1 < 2.

package main

import (
"fmt"
)

type Label int

type Pos struct {
x, y int
}

type DisjointSets struct {
parents map[Pos]Pos
rank    map[Pos]int
labels  map[Pos]Label
}

func NewDisjointSets() *DisjointSets {
return &DisjointSets{
parents: make(map[Pos]Pos),
rank:    make(map[Pos]int),
labels:  make(map[Pos]Label),
}
}

func (s *DisjointSets) MakeSet(u Pos, label Label) {
// ignore if set already exists
if _, ok := s.parents[u]; ok {
return
}

s.labels[u] = label
s.parents[u] = u
s.rank[u] = 1
}

func (s *DisjointSets) FindSet(u Pos) Pos {
assertExists(s.parents, u)

var path []Pos

n := u

for {
if s.parents[n] == n {
break
}

path = append(path, n)
n = s.parents[n]
}

// path compression
for _, p := range path {
s.parents[p] = n
}

return n
}

func (s *DisjointSets) Label(u Pos) Label {
assertExists(s.parents, u)

parent := s.FindSet(u)
return s.labels[parent]
}

func (s *DisjointSets) union(u, v Pos) {
assertExists(s.parents, u, v)

uParent := s.FindSet(u)
vParent := s.FindSet(v)

// already in the same set
if uParent == vParent {
return
}

uLabel := s.Label(u)
vLabel := s.Label(v)

label := minLabel(uLabel, vLabel)

if s.rank[uParent] == s.rank[vParent] {
s.rank[uParent]++
s.parents[vParent] = uParent
s.labels[uParent] = label
} else if s.rank[uParent] > s.rank[vParent] {
s.parents[vParent] = uParent
s.labels[uParent] = label
} else {
s.parents[uParent] = vParent
s.labels[vParent] = label
}
}

func (s *DisjointSets) SameComponent(u, v Pos) bool {
assertExists(s.parents, u, v)

return s.FindSet(u) == s.FindSet(v)
}

func (s *DisjointSets) Dump() {
fmt.Println("====================")

for u, parent := range s.parents {
fmt.Printf(
"(%d,%d) parent: (%d,%d) label %d\n",
u.x,
u.y,
parent.x,
parent.y,
s.Label(u),
)
}
}

func minLabel(labels ...Label) Label {
if len(labels) == 0 {
panic("len(labels) must be > 0")
}

var min = labels[0]

for _, label := range labels {
if label < min {
min = label
}
}

return min
}

func assertExists(parents map[Pos]Pos, check ...Pos) {
for _, u := range check {
if _, ok := parents[u]; !ok {
panic(fmt.Sprintf("Pos (%d,%d) is not a known set.", u.x, u.y))
}
}
}

func main() {
sets := NewDisjointSets()

sets.MakeSet(Pos{0, 0}, 1)
sets.MakeSet(Pos{1, 0}, 2)
sets.MakeSet(Pos{2, 0}, 2)

sets.MakeSet(Pos{0, 1}, 1)
sets.MakeSet(Pos{1, 1}, 3)
sets.MakeSet(Pos{2, 1}, 2)

sets.MakeSet(Pos{0, 2}, 3)
sets.MakeSet(Pos{1, 2}, 3)
sets.MakeSet(Pos{2, 2}, 4)

sets.Dump()

sets.union(Pos{0, 0}, Pos{1, 0})
sets.union(Pos{0, 0}, Pos{1, 1})

sets.Dump()

fmt.Printf("%#v\n", sets.SameComponent(Pos{1, 1}, Pos{2, 2}))
fmt.Printf("%#v\n", sets.SameComponent(Pos{1, 1}, Pos{0, 0}))
fmt.Printf("%#v\n", sets.SameComponent(Pos{0, 2}, Pos{0, 2}))

sets.union(Pos{1, 1}, Pos{1, 1})

sets.Dump()

sets.union(Pos{1, 1}, Pos{2, 2})

sets.Dump()
}

The runnable version can be found here: https://play.golang.org/p/iEF1q2-upJ

Thanks for any feedback on style and functionality.