quick-find.go

// quick-find

type QuickFindSet struct {
    numOfComponents uint
    items []uint
}

func NewSet(n uint) QuickFindSet {
    set := QuickFindSet{ numOfComponents: n, items: make([]uint, n) }
    for i, _ := range set.items { set.items[i] = uint(i) }
    return set
}

func (set *QuickFindSet) Count() uint { return set.numOfComponents }
func (set *QuickFindSet) IsConnected (p, q uint) bool { return set.Find(p) == set.Find(q) }

func (set *QuickFindSet) Find(p uint) uint { return set.items[p] }

func (set *QuickFindSet) Union(p, q uint) {
    rootP := set.Find(p)
    rootQ := set.Find(q)

    if rootP == rootQ { return }

    for i, _ := range set.items {
        if set.items[i] == rootP { set.items[i] = rootQ }
    }

    set.numOfComponents--
}

quick-union.go

// weighted quick-union

type QuickUnionSet struct {
    numOfComponents uint
    items []uint
    sizes []uint
}

func NewSet(n uint) QuickUnionSet {
    set := QuickUnionSet{ numOfComponents: n, items: make([]uint, n), sizes: make([]uint, n) }
    for i, _ := range set.items {
        set.items[i] = uint(i)
        set.sizes[i] = uint(1)
    }
    return set
}

func (set *QuickUnionSet) Count() uint { return set.numOfComponents }
func (set *QuickUnionSet) IsConnected (p, q uint) bool { return set.Find(p) == set.Find(q) }

func (set *QuickUnionSet) Find(p uint) uint {
    for p != set.items[p] { p = set.items[p] }
    return p
}

func (set *QuickUnionSet) Union(p, q uint) {
    rootP := set.Find(p)
    rootQ := set.Find(q)

    if rootP == rootQ { return }

    if set.sizes[rootP] < set.sizes[rootQ] {
        set.items[rootP] = rootQ
        set.sizes[rootQ] += set.sizes[rootP]
    } else {
        set.items[rootQ] = rootP
        set.sizes[rootP] += set.sizes[rootQ]
    }

    set.numOfComponents--
}