btree: add a way to visualize the tree

btree/btree.go

@@ -12,8 +12,8 @@ type data struct {
 
 // Tree is the tree itself
 type Tree struct {
-	root *node // Pointer to the node root
-	t    int   // Minimum degree
+	root   *node // Pointer to the node root
+	degree int   // Minimum degree
 }
 
 type node struct {
@@ -27,11 +27,11 @@ type node struct {
 // Constructors
 
 // NewBtree creates a new btree
-func NewBtree(t int) *Tree {
+func NewBtree(degree int) *Tree {
 
 	return &Tree{
-		root: nil,
-		t:    t,
+		root:   nil,
+		degree: degree,
 	}
 }
 
@@ -62,6 +62,15 @@ func (a *data) gt(b *data) bool {
 
 // Tree methods
 
+// Visualize the tree
+func (t *Tree) Visualize() {
+	if t.root == nil {
+		fmt.Println("The tree is empty")
+		return
+	}
+	t.root.visualize("", true)
+}
+
 // Traverse the tree
 func (t *Tree) Traverse() {
 
@@ -113,8 +122,8 @@ func (t *Tree) Insert(key, value string) {
 
 	// If the tree is empty
 	if t.root == nil {
-		t.root = newNode(t.t, true)
-		t.root.keys[0] = k
+		t.root = newNode(t.degree, true)
+		t.root.insertKey(0, k)
 		t.root.numberOfKeys = 1
 		return
 	}
@@ -127,7 +136,7 @@ func (t *Tree) Insert(key, value string) {
 	}
 
 	// If the root is full, then the tree grows in height
-	s := newNode(t.t, false)
+	s := newNode(t.degree, false)
 
 	// Make the old root as a child of the new root
 	s.children[0] = t.root
@@ -149,6 +158,37 @@ func (t *Tree) Insert(key, value string) {
 
 // node methods
 
+// insertKey at i place in the node
+func (n *node) insertKey(i int, k *data) {
+	n.keys[i] = k
+}
+
+func (n *node) visualize(prefix string, isEnd bool) {
+	fmt.Print(prefix)
+	nextPrefix := "    "
+	if isEnd {
+		fmt.Print("└── ")
+	} else {
+		fmt.Print("├── ")
+		nextPrefix = "│   "
+	}
+	for i := 0; i < n.numberOfKeys-1; i++ {
+		fmt.Printf("%s:%s ", n.keys[i].key, n.keys[i].value)
+	}
+	fmt.Printf("%s:%s", n.keys[n.numberOfKeys-1].key, n.keys[n.numberOfKeys-1].value)
+	fmt.Printf("\n")
+
+	if n.isLeaf {
+		return
+	}
+
+	for i := 0; i < n.numberOfKeys; i++ {
+		n.children[i].visualize(fmt.Sprintf("%s%s", prefix, nextPrefix), false)
+	}
+
+	n.children[n.numberOfKeys].visualize(fmt.Sprintf("%s%s", prefix, nextPrefix), true)
+}
+
 // traverse all nodes in a subtree rooted with this node
 func (n *node) traverse() {
 
@@ -209,7 +249,7 @@ func (n *node) insertNonFull(k *data) {
 		}
 
 		// Insert the new key at the found location
-		n.keys[i+1] = k
+		n.insertKey(i+1, k)
 		n.numberOfKeys++
 		return
 	}

btree/btree_test.go

@@ -106,6 +106,9 @@ func ExampleTree_Traverse() {
 		tree.Insert(k[0], k[1])
 	}
 
+	tree.Visualize()
+	fmt.Println("")
+
 	tree.Traverse()
 	fmt.Println("")
 
@@ -133,7 +136,14 @@ func ExampleTree_Traverse() {
 	tree.Traverse()
 
 	// Output:
-	// a:A b:B c:C d:D e:E f:F g:G h:H i:I j:J k:K l:L m:M n:N o:O p:P q:Q
+	// └── c:C f:F i:I l:L
+	//     ├── a:A b:B
+	//     ├── d:D e:E
+	//     ├── g:G h:H
+	//     ├── j:J k:K
+	//     └── m:M n:N o:O p:P q:Q
+	//
+	//  a:A b:B c:C d:D e:E f:F g:G h:H i:I j:J k:K l:L m:M n:N o:O p:P q:Q
 	//  a:A b:B c:C d:D f:F g:G h:H i:I j:J k:K l:L m:M n:N o:O p:P q:Q
 	//  a:A b:B c:C d:D f:F g:G h:H i:I k:K l:L m:M n:N o:O p:P q:Q
 	//  b:B c:C d:D f:F g:G h:H i:I k:K l:L m:M n:N o:O p:P q:Q

main.go

@@ -11,16 +11,28 @@ func main() {
 	// Example program
 	t := btree.NewBtree(3)
 	t.Insert("animal", "dog")
+
+	s, _ := t.Search("animal")
+	fmt.Println(s)
+
 	t.Insert("potato", "fries")
 	t.Insert("6", "number")
 	t.Insert("12", "number")
 	t.Insert("car", "ferrari")
 	t.Insert("7", "number")
 	t.Insert("plane", "airbus")
+	t.Insert("animal", "cat")
+
+	s, _ = t.Search("animal")
+	fmt.Println(s)
 
 	fmt.Print("Traversal of the constructed tree is")
 	t.Traverse()
 
+	fmt.Println("")
+
+	t.Visualize()
+
 	if val, err := t.Search("6"); err != nil {
 		fmt.Print("\nNot Present")
 	} else {
@@ -32,4 +44,5 @@ func main() {
 	} else {
 		fmt.Printf("\nPresent: value '%s'", val)
 	}
+
 }