Union Find

quick_find.go

package main

import "fmt"

type quickFind []int

func NewQuickFind(n int) quickFind {
  u := make(quickFind, n)
	for i := 0; i < n; i++ {
		u[i] = i
	}
	return u
}

func (u quickFind) union(p, q int) {
	pid, qid := u[p], u[q]
	for i, v := range u {
		if v == pid {
			u[i] = qid
		}
	}
}

func (u quickFind) connected(p, q int) bool {
	return u[p] == u[q]
}

func main() {
	u := NewQuickFind(10)

	u.union(1, 9)
	u.union(3, 4)
	u.union(5, 7)
	u.union(9, 2)
	u.union(4, 7)

	cases := [][2]int{
		{1, 2},
		{3, 8},
		{5, 4},
		{5, 6},
	}

	for _, c := range cases {
		if u.connected(c[0], c[1]) {
			fmt.Printf("%d --> %d\n", c[0], c[1])
		} else {
			fmt.Printf("%d -/> %d\n", c[0], c[1])
		}
	}
}

quick_union.go

package main

import "fmt"

type quickUnion []int

func (u quickUnion) root(i int) int {
	for u[i] != i {
		i = u[i]
	}
	return i
}

func (u quickUnion) connected(p, q int) bool {
	return u.root(p) == u.root(q)
}

func (u quickUnion) union(p, q int) {
	i := u.root(p)
	u[i] = u.root(q)
}

func NewQuickUnion(n int) quickUnion {
	u := make(quickUnion, n)
	for i := 0; i < n; i++ {
		u[i] = i
	}
	return u
}

func main() {
	u := NewQuickUnion(10)

	u.union(1, 8)
	u.union(8, 2)
	u.union(7, 5)
	u.union(5, 2)

	cases := [][2]int{
		{1, 2},
		{2, 3},
		{1, 1},
		{5, 1},
	}

	for _, c := range cases {
		if u.connected(c[0], c[1]) {
			fmt.Printf("%d --> %d\n", c[0], c[1])
		} else {
			fmt.Printf("%d -/> %d\n", c[0], c[1])
		}
	}
}

weighted_quick_union_with_path_compression.go

package main

import "fmt"

type quickUnion struct {
	id []int
	sz []int
}

func (u *quickUnion) root(i int) int {
	for u.id[i] != i {
		// path compression
		u.id[i] = u.id[u.id[i]]
		i = u.id[i]
	}
	return i
}

func (u *quickUnion) connected(p, q int) bool {
	return u.root(q) == u.root(p)
}

func (u *quickUnion) union(p, q int) {
	pid := u.root(p)
	qid := u.root(q)
	
	// make the smaller tree a child of the larger one
	if u.sz[pid] > u.sz[qid] {
		u.id[qid] = pid
		
		// the smaller tree's current root will never be used again.
		// so we only need to update the parent's size.
		u.sz[pid] += u.sz[qid]
	} else {
		u.id[pid] = qid
		u.sz[qid] += u.sz[pid]
	}

}

func NewQuickUnion(n int) *quickUnion {
	u := quickUnion{
		id: make([]int, n),
		sz: make([]int, n),
	}
	for i := 0; i < n; i++ {
		u.id[i] = i
		u.sz[i] = 1
	}
	return &u
}

func main() {
	u := NewQuickUnion(10)

	u.union(1, 8)
	u.union(8, 2)
	u.union(7, 5)
	u.union(5, 2)

	cases := [][2]int{
		{1, 2},
		{2, 3},
		{1, 1},
		{5, 1},
	}

	for _, c := range cases {
		if u.connected(c[0], c[1]) {
			fmt.Printf("%d --> %d\n", c[0], c[1])
		} else {
			fmt.Printf("%d -/> %d\n", c[0], c[1])
		}
	}
}