- debugging removal in red black tree (will carry on later)

trees/redblacktree/redblacktree.go

@@ -18,10 +18,12 @@ with this distribution for more information.
 
 // Implementation of Red-black tree.
 // Used by TreeSet and TreeMap.
+// Structure is not thread safe.
 // References: http://en.wikipedia.org/wiki/Red%E2%80%93black_tree
 package redblacktree
 
 import (
+	"fmt"
 	"github.com/emirpasic/gods/utils"
 )
 
@@ -121,18 +123,20 @@ func (tree *Tree) Remove(key interface{}) {
 		node.value = pred.value
 		node = pred
 	}
-	if node.right == nil {
-		child = node.left
-	} else {
-		child = node.right
-	}
-	if node.color == BLACK {
-		node.color = child.color
-		tree.deleteCase1(node)
-	}
-	tree.replaceNode(node, child)
-	if node.parent == nil && child != nil {
-		child.color = BLACK
+	if node.left == nil || node.right == nil {
+		if node.right == nil {
+			child = node.left
+		} else {
+			child = node.right
+		}
+		if node.color == BLACK {
+			node.color = child.color
+			tree.deleteCase1(node)
+		}
+		tree.replaceNode(node, child)
+		if node.parent == nil && child != nil {
+			child.color = BLACK
+		}
 	}
 }
 
@@ -141,6 +145,39 @@ func (tree *Tree) IsEmpty() bool {
 	return tree.root == nil
 }
 
+func (tree *Tree) String() string {
+	if !tree.IsEmpty() {
+		levels := make(map[int][]*Node)
+		lastLevel := 0
+		levels[lastLevel] = append(levels[lastLevel], tree.root)
+		for len(levels[lastLevel]) > 0 {
+			for _, node := range levels[lastLevel] {
+				if node.left != nil {
+					levels[lastLevel+1] = append(levels[lastLevel+1], node.left)
+				}
+				if node.right != nil {
+					levels[lastLevel+1] = append(levels[lastLevel+1], node.right)
+				}
+			}
+			lastLevel += 1
+		}
+		str := "tree\n"
+
+		for i := 0; i < lastLevel; i++ {
+			for _, node := range levels[i] {
+				str +=  fmt.Sprintf("%s", node)
+			}
+			str += "\n"
+		}
+		return str
+	}
+	return "tree\n"
+}
+
+func (node *Node) String() string {
+	return fmt.Sprintf("(%v)", node.key)
+}
+
 func (tree *Tree) lookup(key interface{}) *Node {
 	node := tree.root
 	for node != nil {
@@ -347,10 +384,10 @@ func (tree *Tree) deleteCase5(node *Node) {
 func (tree *Tree) deleteCase6(node *Node) {
 	node.sibling().color = node.parent.color
 	node.parent.color = BLACK
-	if node == node.parent.left {
+	if node == node.parent.left && node.sibling().right.color == RED {
 		node.sibling().right.color = BLACK
 		tree.rotateLeft(node.parent)
-	} else {
+	} else if node.sibling().left.color == RED {
 		node.sibling().left.color = BLACK
 		tree.rotateRight(node.parent)
 	}

trees/redblacktree/redblacktree_test.go

@@ -2,13 +2,17 @@ package redblacktree
 
 import (
 	"testing"
+	"log"
 )
 
-func TestPutGet(t *testing.T) {
+func TestRedBlackTree(t *testing.T) {
 
 	tree := NewWithIntComparator()
 
+	// insertions
 	tree.Put(5, "e")
+	tree.Put(6, "f")
+	tree.Put(7, "g")
 	tree.Put(3, "c")
 	tree.Put(4, "d")
 	tree.Put(1, "x")
@@ -16,18 +20,51 @@ func TestPutGet(t *testing.T) {
 	tree.Put(1, "a") //overwrite
 
 	// key,expectedValue,expectedFound
-	tests := [][]interface{}{
+	tests1 := [][]interface{}{
+		{1, "a", true},
+		{2, "b", true},
+		{3, "c", true},
+		{4, "d", true},
+		{5, "e", true},
+		{6, "f", true},
+		{7, "g", true},
+		{8, nil, false},
+	}
+
+	for _, test := range tests1 {
+		// retrievals
+		actualValue, actualFound := tree.Get(test[0])
+		if actualValue != test[1] || actualFound != test[2] {
+			t.Errorf("Got %v expected %v", actualValue, test[1])
+		}
+	}
+
+	// removals
+	log.Println(tree)
+	tree.Remove(5)
+	log.Println(tree)
+
+	tree.Remove(6)
+	tree.Remove(7)
+	tree.Remove(8)
+
+	tests2 := [][]interface{}{
 		{1, "a", true},
 		{2, "b", true},
 		{3, "c", true},
 		{4, "d", true},
 		{5, nil, false},
+		{6, nil, false},
+		{7, nil, false},
+		{8, nil, false},
 	}
 
-	for _, test := range tests {
+	for _, test := range tests2 {
+		// retrievals
 		actualValue, actualFound := tree.Get(test[0])
 		if actualValue != test[1] || actualFound != test[2] {
 			t.Errorf("Got %v expected %v", actualValue, test[1])
 		}
 	}
+
 }