- btree Height() and String() implementations

trees/btree/btree.go

@@ -10,9 +10,11 @@
 package btree
 
 import (
+	"bytes"
 	"fmt"
 	"github.com/emirpasic/gods/trees"
 	"github.com/emirpasic/gods/utils"
+	"strings"
 )
 
 func assertTreeImplementation() {
@@ -114,23 +116,49 @@ func (tree *Tree) Clear() {
 	tree.size = 0
 }
 
-// String returns a string representation of container
+// Height returns the height of the tree.
+func (tree *Tree) Height() int {
+	return tree.Root.height()
+}
+
+// String returns a string representation of container (for debugging purposes)
 func (tree *Tree) String() string {
-	str := "BTree\n"
+	var buffer bytes.Buffer
+	buffer.WriteString("BTree\n")
 	if !tree.Empty() {
-		str += tree.Root.String()
+		tree.output(&buffer, tree.Root, 0, true)
 	}
-	return str
-}
-
-func (node *Node) String() string {
-	return fmt.Sprintf("%v", node.Entries)
+	return buffer.String()
 }
 
 func (entry *Entry) String() string {
 	return fmt.Sprintf("%v", entry.Key)
 }
 
+func (tree *Tree) output(buffer *bytes.Buffer, node *Node, level int, isTail bool) {
+	for e := 0; e < len(node.Entries)+1; e++ {
+		if e < len(node.Children) {
+			buffer.WriteString(strings.Repeat("   ", level))
+			tree.output(buffer, node.Children[e], level+1, true)
+		}
+		if e < len(node.Entries) {
+			buffer.WriteString(strings.Repeat("   ", level))
+			buffer.WriteString(fmt.Sprintf("%v", node.Entries[e].Key) + "\n")
+		}
+	}
+}
+
+func (node *Node) height() int {
+	height := 0
+	for ; node != nil; node = node.Children[0] {
+		height++
+		if len(node.Children) == 0 {
+			break
+		}
+	}
+	return height
+}
+
 func (tree *Tree) isLeaf(node *Node) bool {
 	return len(node.Children) == 0
 }
@@ -183,7 +211,6 @@ func (tree *Tree) insert(node *Node, entry *Entry) (inserted bool) {
 func (tree *Tree) insertIntoLeaf(node *Node, entry *Entry) (inserted bool) {
 	insertPosition, found := tree.search(node, entry)
 	if found {
-		node.Entries[insertPosition] = nil // GC
 		node.Entries[insertPosition] = entry
 		return false
 	}
@@ -197,7 +224,6 @@ func (tree *Tree) insertIntoLeaf(node *Node, entry *Entry) (inserted bool) {
 func (tree *Tree) insertIntoInternal(node *Node, entry *Entry) (inserted bool) {
 	insertPosition, found := tree.search(node, entry)
 	if found {
-		node.Entries[insertPosition] = nil // GC
 		node.Entries[insertPosition] = entry
 		return false
 	}
@@ -240,8 +266,6 @@ func (tree *Tree) splitNonRoot(node *Node) {
 	copy(parent.Children[insertPosition+2:], parent.Children[insertPosition+1:])
 	parent.Children[insertPosition+1] = right
 
-	node = nil // GC
-
 	tree.split(parent)
 }
 

trees/btree/btree_test.go

@@ -192,6 +192,48 @@ func TestBTree_insert3(t *testing.T) {
 	assertValidTreeNode(t, tree.Root.Children[2], 3, 0, []int{70, 80, 90})
 }
 
+func TestBTree_height(t *testing.T) {
+	tree := NewWithIntComparator(3)
+	if actualValue, expectedValue := tree.Height(), 0; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(1, 0)
+	if actualValue, expectedValue := tree.Height(), 1; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(2, 1)
+	if actualValue, expectedValue := tree.Height(), 1; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(3, 2)
+	if actualValue, expectedValue := tree.Height(), 2; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(4, 2)
+	if actualValue, expectedValue := tree.Height(), 2; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(5, 2)
+	if actualValue, expectedValue := tree.Height(), 2; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(6, 2)
+	if actualValue, expectedValue := tree.Height(), 2; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+
+	tree.Put(7, 2)
+	if actualValue, expectedValue := tree.Height(), 3; actualValue != expectedValue {
+		t.Errorf("Got %v expected %v", actualValue, expectedValue)
+	}
+}
+
 func assertValidTree(t *testing.T, tree *Tree, expectedSize int) {
 	if actualValue, expectedValue := tree.size, expectedSize; actualValue != expectedValue {
 		t.Errorf("Got %v expected %v for tree size", actualValue, expectedValue)

trees/redblacktree/redblacktree.go

@@ -268,7 +268,6 @@ func (node *Node) String() string {
 	return fmt.Sprintf("%v", node.Key)
 }
 
-// String returns a string representation of container
 func output(node *Node, prefix string, isTail bool, str *string) {
 	if node.Right != nil {
 		newPrefix := prefix