From a2b057ed414221634aa22d298f0ffe1f1cac6443 Mon Sep 17 00:00:00 2001
From: Emir Pasic <emirpasic@hotmail.com>
Date: Thu, 5 Mar 2015 14:37:28 +0100
Subject: [PATCH] red-black tree, get keys and values in-order (with tests)

---
 .../linkedliststack/linkedliststack_test.go   |  3 --
 trees/redblacktree/redblacktree.go            | 47 +++++++++++++++++
 trees/redblacktree/redblacktree_test.go       | 52 ++++++++++++++++++-
 3 files changed, 98 insertions(+), 4 deletions(-)

diff --git a/stacks/linkedliststack/linkedliststack_test.go b/stacks/linkedliststack/linkedliststack_test.go
index 6e7f115..9503736 100644
--- a/stacks/linkedliststack/linkedliststack_test.go
+++ b/stacks/linkedliststack/linkedliststack_test.go
@@ -19,7 +19,6 @@ with this distribution for more information.
 package linkedliststack
 
 import (
-	"fmt"
 	"testing"
 )
 
@@ -36,8 +35,6 @@ func TestLinkedListStack(t *testing.T) {
 	stack.Push(2)
 	stack.Push(3)
 
-	fmt.Println(stack)
-
 	if actualValue := stack.Empty(); actualValue != false {
 		t.Errorf("Got %v expected %v", actualValue, false)
 	}
diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go
index 9e8baff..e54c1f0 100644
--- a/trees/redblacktree/redblacktree.go
+++ b/trees/redblacktree/redblacktree.go
@@ -25,6 +25,7 @@ package redblacktree
 
 import (
 	"fmt"
+	"github.com/emirpasic/gods/stacks/linkedliststack"
 	"github.com/emirpasic/gods/utils"
 )
 
@@ -154,6 +155,24 @@ func (tree *Tree) Size() int {
 	return tree.size
 }
 
+// Returns all keys in-order
+func (tree *Tree) Keys() []interface{} {
+	keys := make([]interface{}, tree.size)
+	for i, node := range tree.inOrder() {
+		keys[i] = node.key
+	}
+	return keys
+}
+
+// Returns all values in-order based on the key
+func (tree *Tree) Values() []interface{} {
+	values := make([]interface{}, tree.size)
+	for i, node := range tree.inOrder() {
+		values[i] = node.value
+	}
+	return values
+}
+
 func (tree *Tree) String() string {
 	str := "RedBlackTree\n"
 	if !tree.Empty() {
@@ -166,6 +185,34 @@ func (node *node) String() string {
 	return fmt.Sprintf("%v", node.key)
 }
 
+// Returns all nodes in order
+func (tree *Tree) inOrder() []*node {
+	nodes := make([]*node, tree.size)
+	if tree.size > 0 {
+		current := tree.root
+		stack := linkedliststack.New()
+		done := false
+		count := 0
+		for !done {
+			if current != nil {
+				stack.Push(current)
+				current = current.left
+			} else {
+				if !stack.Empty() {
+					currentPop, _ := stack.Pop()
+					current = currentPop.(*node)
+					nodes[count] = current
+					count += 1
+					current = current.right
+				} else {
+					done = true
+				}
+			}
+		}
+	}
+	return nodes
+}
+
 func output(node *node, prefix string, isTail bool, str *string) {
 	if node.right != nil {
 		newPrefix := prefix
diff --git a/trees/redblacktree/redblacktree_test.go b/trees/redblacktree/redblacktree_test.go
index 2685768..d3e9341 100644
--- a/trees/redblacktree/redblacktree_test.go
+++ b/trees/redblacktree/redblacktree_test.go
@@ -19,6 +19,7 @@ with this distribution for more information.
 package redblacktree
 
 import (
+	"fmt"
 	"testing"
 )
 
@@ -36,12 +37,22 @@ func TestRedBlackTree(t *testing.T) {
 	tree.Put(2, "b")
 	tree.Put(1, "a") //overwrite
 
+	// Test Size()
 	if actualValue := tree.Size(); actualValue != 7 {
 		t.Errorf("Got %v expected %v", actualValue, 7)
 	}
 
-	// key,expectedValue,expectedFound
+	// test Keys()
+	if actualValue, expactedValue := fmt.Sprintf("%d%d%d%d%d%d%d", tree.Keys()...), "1234567"; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
 
+	// test Values()
+	if actualValue, expactedValue := fmt.Sprintf("%s%s%s%s%s%s%s", tree.Values()...), "abcdefg"; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
+
+	// key,expectedValue,expectedFound
 	tests1 := [][]interface{}{
 		{1, "a", true},
 		{2, "b", true},
@@ -68,6 +79,17 @@ func TestRedBlackTree(t *testing.T) {
 	tree.Remove(8)
 	tree.Remove(5)
 
+	// Test Keys()
+	if actualValue, expactedValue := fmt.Sprintf("%d%d%d%d", tree.Keys()...), "1234"; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
+
+	// test Values()
+	if actualValue, expactedValue := fmt.Sprintf("%s%s%s%s", tree.Values()...), "abcd"; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
+
+	// Test Size()
 	if actualValue := tree.Size(); actualValue != 4 {
 		t.Errorf("Got %v expected %v", actualValue, 7)
 	}
@@ -91,4 +113,32 @@ func TestRedBlackTree(t *testing.T) {
 		}
 	}
 
+	// removals
+	tree.Remove(1)
+	tree.Remove(4)
+	tree.Remove(2)
+	tree.Remove(3)
+	tree.Remove(2)
+	tree.Remove(2)
+
+	// Test Keys()
+	if actualValue, expactedValue := fmt.Sprintf("", tree.Keys()...), ""; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
+
+	// test Values()
+	if actualValue, expactedValue := fmt.Sprintf("", tree.Values()...), ""; actualValue != expactedValue {
+		t.Errorf("Got %v expected %v", actualValue, expactedValue)
+	}
+
+	// Test Size()
+	if actualValue := tree.Size(); actualValue != 0 {
+		t.Errorf("Got %v expected %v", actualValue, 0)
+	}
+
+	// Test Empty()
+	if actualValue := tree.Empty(); actualValue != true {
+		t.Errorf("Got %v expected %v", actualValue, true)
+	}
+
 }