Merge pull request #338 from peterstace/disjointset

Add DisjointSet data structure

geom/alg_disjoint_set.go

@@ -0,0 +1,55 @@
+package geom
+
+// disjointSet implements a standard disjoint-set data structure (also known as
+// a union-find structure or a merge-find structure). It stores a collection of
+// disjoint (non-overlapping) sets. The set elements are integers (which may be
+// mapped externally to more complicated types).
+type disjointSet struct {
+	parent []int // self reference indicates root
+	rank   []int
+}
+
+// newDistointSet creates a new disjoint set containing n sets, each with a
+// single item. The items are 0 (inclusive) through to n (exclusive).
+func newDisjointSet(n int) disjointSet {
+	set := disjointSet{make([]int, n), make([]int, n)}
+	for i := range set.parent {
+		set.parent[i] = i
+	}
+	return set
+}
+
+// find searches for the representative for the set containing x. All elements
+// of the set containing x will have the same representative. To find out if
+// two elements are in the same set, find can be used on each element and the
+// representatives compared.
+func (s disjointSet) find(x int) int {
+	root := x
+	for s.parent[root] != root {
+		root = s.parent[root]
+	}
+	for s.parent[x] != root {
+		parent := s.parent[x]
+		s.parent[x] = root
+		x = parent
+	}
+	return root
+}
+
+// union merges the set containing x with the set containing y.
+func (s disjointSet) union(x, y int) {
+	x = s.find(x)
+	y = s.find(y)
+	if x == y {
+		return
+	}
+
+	if s.rank[x] < s.rank[y] {
+		x, y = y, x
+	}
+
+	s.parent[y] = x
+	if s.rank[x] == s.rank[y] {
+		s.rank[x]++
+	}
+}

geom/alg_disjoint_set_test.go

@@ -0,0 +1,111 @@
+package geom
+
+import (
+	"strconv"
+	"testing"
+)
+
+func TestDisjointSet(t *testing.T) {
+	type intPair struct{ a, b int }
+	for idx, tc := range []struct {
+		size   int
+		merges []intPair
+	}{
+		{
+			size:   0,
+			merges: nil,
+		},
+		{
+			size:   1,
+			merges: []intPair{{0, 0}},
+		},
+		{
+			size:   2,
+			merges: []intPair{{0, 1}},
+		},
+		{
+			size:   2,
+			merges: []intPair{{1, 0}},
+		},
+		{
+			size:   2,
+			merges: []intPair{{1, 1}},
+		},
+		{
+			size: 3,
+			merges: []intPair{
+				{1, 0},
+				{0, 2},
+			},
+		},
+		{
+			size: 4,
+			merges: []intPair{
+				{2, 1},
+				{0, 3},
+				{1, 0},
+			},
+		},
+		{
+			size: 5,
+			merges: []intPair{
+				{2, 1},
+				{0, 3},
+				{1, 0},
+				{3, 4},
+			},
+		},
+	} {
+		t.Run(strconv.Itoa(idx), func(t *testing.T) {
+			t.Logf("num elements: %d", tc.size)
+			simpleSet := newSimpleDisjointSet(tc.size)
+			fastSet := newDisjointSet(tc.size)
+			for _, m := range tc.merges {
+				t.Logf("merging %d and %d", m.a, m.b)
+				fastSet.union(m.a, m.b)
+				simpleSet.union(m.a, m.b)
+				for i := 0; i < tc.size; i++ {
+					for j := 0; j < tc.size; j++ {
+						gotFast := fastSet.find(i) == fastSet.find(j)
+						gotSimple := simpleSet.find(i) == simpleSet.find(j)
+						if gotFast != gotSimple {
+							t.Errorf("mismatch between %d and %d in same set: fast=%v simple=%v", i, j, gotFast, gotSimple)
+						}
+					}
+				}
+			}
+		})
+	}
+}
+
+// simpleDisjointSet is a _simple_ but _inefficient_ implementation
+// of a disjoint set data structure. It's used as a reference
+// implementation for testing.
+type simpleDisjointSet struct {
+	// For the simple implementation, we store the set identifier
+	// for each element directly. This results in very simple
+	// operations, but linear union time complexity.
+	set []int
+}
+
+func newSimpleDisjointSet(size int) simpleDisjointSet {
+	items := make([]int, size)
+	for i := range items {
+		items[i] = i
+	}
+	return simpleDisjointSet{items}
+}
+
+func (s simpleDisjointSet) find(x int) int {
+	return s.set[x]
+}
+
+func (s simpleDisjointSet) union(x, y int) {
+	setX := s.find(x)
+	setY := s.find(y)
+	for i := range s.set {
+		if s.set[i] == setY {
+			s.set[i] = setX
+		}
+	}
+}