1
\$\begingroup\$

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.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.