Add Count() method to the API

README.md

@@ -19,6 +19,7 @@ uf.Union(4, 6)
 fmt.PrintLn(uf.Find(2)) // Prints 1
 fmt.PrintLn(uf.Connected(1, 2)) // Prints true
 fmt.PrintLn(uf.Connected(1, 3)) // Prints false
+fmt.PrintLn(uf.Count()) // Prints 6
 
 ```
 
@@ -39,8 +40,11 @@ Same as Find.
 #### `uf.Connected(p, q)`
 Checks if p and q are connected.
 
+#### `uf.Count()`
+Returns the number of independent components.
+
 ## EXTRA
 There is also a goroutine safe version of unionfind in the file `safe-unionfind`.
 
 ## Licence
-MIT - Theo Despoudis
+MIT - Theo Despoudis

safe-unionfind.go

@@ -27,7 +27,6 @@ SOFTWARE.
 
 // Thread safe version of Union-Find using plain locks.
 
-
 type ThreadSafeUnionFind struct {
 	sync.RWMutex
 	uf *UnionFind
@@ -47,7 +46,6 @@ func (suf *ThreadSafeUnionFind) Union(p int, q int) {
 	suf.uf.Union(p, q)
 }
 
-
 func (suf *ThreadSafeUnionFind) Root(p int) int {
 	suf.Lock()
 	defer suf.Unlock()
@@ -68,4 +66,9 @@ func (suf *ThreadSafeUnionFind) Connected(p int, q int) bool {
 	return suf.uf.Root(p) == suf.uf.Root(p)
 }
 
+func (suf *ThreadSafeUnionFind) Count() int {
+	suf.Lock()
+	defer suf.Unlock()
 
+	return suf.uf.Count()
+}

unionfind.go

@@ -36,8 +36,9 @@ SOFTWARE.
 
 //
 type UnionFind struct {
-	root []int
-	size []int
+	root  []int
+	size  []int
+	count int
 }
 
 // New returns an initialized list of size
@@ -50,6 +51,7 @@ func (uf *UnionFind) init(size int) *UnionFind {
 	uf = new(UnionFind)
 	uf.root = make([]int, size)
 	uf.size = make([]int, size)
+	uf.count = size
 
 	for i := 0; i < size; i++ {
 		uf.root[i] = i
@@ -61,9 +63,17 @@ func (uf *UnionFind) init(size int) *UnionFind {
 
 // Union connects p and q by finding their roots and comparing their respective
 // size arrays to keep the tree flat
+// Time complexity is amortized O(1)
 func (uf *UnionFind) Union(p int, q int) {
+	if p == q {
+		return
+	}
+
 	qRoot := uf.Root(q)
 	pRoot := uf.Root(p)
+	if qRoot == pRoot {
+		return
+	}
 
 	if uf.size[qRoot] < uf.size[pRoot] {
 		uf.root[qRoot] = uf.root[pRoot]
@@ -72,11 +82,14 @@ func (uf *UnionFind) Union(p int, q int) {
 		uf.root[pRoot] = uf.root[qRoot]
 		uf.size[qRoot] += uf.size[pRoot]
 	}
+
+	uf.count--
 }
 
 // Root or Find traverses each parent element while compressing the
 // levels to find the root element of p
 // If we attempt to access an element outside the array it returns -1
+// Time complexity is amortized O(1)
 func (uf *UnionFind) Root(p int) int {
 	if p > len(uf.root)-1 {
 		return -1
@@ -99,3 +112,9 @@ func (uf *UnionFind) Find(p int) int {
 func (uf *UnionFind) Connected(p int, q int) bool {
 	return uf.Root(p) == uf.Root(q)
 }
+
+// Count returns the number of independent items
+// Time complexity is O(1)
+func (uf *UnionFind) Count() int {
+	return uf.count
+}

unionfind_test.go

@@ -19,6 +19,7 @@ var _ = Suite(&MySuite{})
 func (s *MySuite) TestNew(c *C) {
 	uf := New(10)
 
+	c.Assert(uf.count, Equals, 10)
 	c.Assert(uf.size, DeepEquals, []int{1, 1, 1, 1, 1, 1, 1, 1, 1, 1})
 	c.Assert(uf.root, DeepEquals, []int{0, 1, 2, 3, 4, 5, 6, 7, 8, 9})
 }
@@ -29,9 +30,11 @@ func (s *MySuite) TestRootEmpty(c *C) {
 
 	c.Assert(uf.Root(10), Equals, -1)
 	c.Assert(uf.Root(9), Equals, 9)
+	c.Assert(uf.Find(9), Equals, 9)
+	c.Assert(uf.Count(), Equals, 10)
 }
 
-// Test Root when It has connected components
+// Test Root when it has connected components
 func (s *MySuite) TestRootConnected(c *C) {
 	uf := New(10)
 
@@ -43,6 +46,7 @@ func (s *MySuite) TestRootConnected(c *C) {
 	c.Assert(uf.Root(8), Equals, 2)
 	c.Assert(uf.Root(2), Equals, 2)
 	c.Assert(uf.Root(5), Equals, 6)
+	c.Assert(uf.Count(), Equals, 6)
 }
 
 // Test Union when when we connect the same item
@@ -52,6 +56,20 @@ func (s *MySuite) TestUnionSame(c *C) {
 	uf.Union(2, 2)
 
 	c.Assert(uf.Root(2), Equals, 2)
+	c.Assert(uf.Count(), Equals, 10)
+}
+
+// Test Union when when we connect items with the same root
+func (s *MySuite) TestUnionSameRoot(c *C) {
+	uf := New(5)
+	uf.Union(0, 1)
+	uf.Union(1, 2)
+	uf.Union(2, 3)
+	c.Assert(uf.Count(), Equals, 2)
+
+	uf.Union(0, 3)
+
+	c.Assert(uf.Count(), Equals, 2)
 }
 
 // Test Union checks size array when it assigns the root
@@ -63,4 +81,6 @@ func (s *MySuite) TestUnionChecksSizes(c *C) {
 	uf.Union(2, 3)
 
 	c.Assert(uf.size, DeepEquals, []int{1, 4, 1, 1, 1})
+	c.Assert(uf.Connected(0, 3), Equals, true)
+	c.Assert(uf.Count(), Equals, 2)
 }