- adding save color(node) function to guard against nil-nodes

trees/redblacktree/redblacktree.go

@@ -130,7 +130,7 @@ func (tree *Tree) Remove(key interface{}) {
 			child = node.right
 		}
 		if node.color == BLACK {
-			node.color = child.color
+			node.color = color(child)
 			tree.deleteCase1(node)
 		}
 		tree.replaceNode(node, child)
@@ -146,36 +146,43 @@ func (tree *Tree) IsEmpty() bool {
 }
 
 func (tree *Tree) String() string {
+	str := "RedBlackTree\n"
 	if !tree.IsEmpty() {
-		levels := make(map[int][]*Node)
+		output(tree.root, "", true, &str)
-		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)
 	}
+	return str
+}
+
+func (node *Node) String() string {
+	return fmt.Sprintf("%v", node.key)
+}
+
+func output(node *Node, prefix string, isTail bool, str *string) {
 	if node.right != nil {
-					levels[lastLevel+1] = append(levels[lastLevel+1], node.right)
+		newPrefix := prefix
+		if isTail {
+			newPrefix += "│   "
+		} else {
+			newPrefix += "    "
 		}
+		output(node.right, newPrefix, false, str)
 	}
-			lastLevel += 1
+	*str += prefix
+	if isTail {
+		*str += "└── "
+	} else {
+		*str += "┌── "
 	}
-		str := "tree\n"
+	*str += node.String() + "\n"
-
+	if node.left != nil {
-		for i := 0; i < lastLevel; i++ {
+		newPrefix := prefix
-			for _, node := range levels[i] {
+		if isTail {
-				str +=  fmt.Sprintf("%s", node)
+			newPrefix += "    "
+		} else {
+			newPrefix += "│   "
 		}
-			str += "\n"
+		output(node.left, newPrefix, true, str)
 	}
-		return str
-	}
-	return "tree\n"
-}
-
-func (node *Node) String() string {
-	return fmt.Sprintf("(%v)", node.key)
 }
 
 func (tree *Tree) lookup(key interface{}) *Node {
@@ -265,18 +272,15 @@ func (tree *Tree) insertCase1(node *Node) {
 }
 
 func (tree *Tree) insertCase2(node *Node) {
-	if node.parent.color == BLACK {
+	if color(node.parent) == BLACK {
-		return /* Tree is still valid */
+		return
 	}
 	tree.insertCase3(node)
 }
 
 func (tree *Tree) insertCase3(node *Node) {
 	uncle := node.uncle()
-	if uncle == nil {
+	if color(uncle) == RED {
-		return
-	}
-	if uncle.color == RED {
 		node.parent.color = BLACK
 		uncle.color = BLACK
 		node.grandparent().color = RED
@@ -287,10 +291,11 @@ func (tree *Tree) insertCase3(node *Node) {
 }
 
 func (tree *Tree) insertCase4(node *Node) {
-	if node == node.parent.right && node.parent == node.grandparent().left {
+	grandparent := node.grandparent()
+	if node == node.parent.right && node.parent == grandparent.left {
 		tree.rotateLeft(node.parent)
 		node = node.left
-	} else if node == node.parent.left && node.parent == node.grandparent().right {
+	} else if node == node.parent.left && node.parent == grandparent.right {
 		tree.rotateRight(node.parent)
 		node = node.right
 	}
@@ -299,11 +304,12 @@ func (tree *Tree) insertCase4(node *Node) {
 
 func (tree *Tree) insertCase5(node *Node) {
 	node.parent.color = BLACK
-	node.grandparent().color = RED
+	grandparent := node.grandparent()
-	if node == node.parent.left && node.parent == node.grandparent().left {
+	grandparent.color = RED
-		tree.rotateRight(node.grandparent())
+	if node == node.parent.left && node.parent == grandparent.left {
-	} else if node == node.parent.right && node.parent == node.grandparent().right {
+		tree.rotateRight(grandparent)
-		tree.rotateLeft(node.grandparent())
+	} else if node == node.parent.right && node.parent == grandparent.right {
+		tree.rotateLeft(grandparent)
 	}
 }
 
@@ -326,9 +332,10 @@ func (tree *Tree) deleteCase1(node *Node) {
 }
 
 func (tree *Tree) deleteCase2(node *Node) {
-	if node.sibling().color == RED {
+	sibling := node.sibling()
+	if color(sibling) == RED {
 		node.parent.color = RED
-		node.sibling().color = BLACK
+		sibling.color = BLACK
 		if node == node.parent.left {
 			tree.rotateLeft(node.parent)
 		} else {
@@ -339,11 +346,12 @@ func (tree *Tree) deleteCase2(node *Node) {
 }
 
 func (tree *Tree) deleteCase3(node *Node) {
-	if node.parent.color == BLACK &&
+	sibling := node.sibling()
-		node.sibling().color == BLACK &&
+	if color(node.parent) == BLACK &&
-		node.sibling().left.color == BLACK &&
+		color(sibling) == BLACK &&
-		node.sibling().right.color == BLACK {
+		color(sibling.left) == BLACK &&
-		node.sibling().color = RED
+		color(sibling.right) == BLACK {
+		sibling.color = RED
 		tree.deleteCase1(node.parent)
 	} else {
 		tree.deleteCase4(node)
@@ -351,11 +359,12 @@ func (tree *Tree) deleteCase3(node *Node) {
 }
 
 func (tree *Tree) deleteCase4(node *Node) {
-	if node.parent.color == RED &&
+	sibling := node.sibling()
-		node.sibling().color == BLACK &&
+	if color(node.parent) == RED &&
-		node.sibling().left.color == BLACK &&
+		color(sibling) == BLACK &&
-		node.sibling().right.color == BLACK {
+		color(sibling.left) == BLACK &&
-		node.sibling().color = RED
+		color(sibling.right) == BLACK {
+		sibling.color = RED
 		node.parent.color = BLACK
 	} else {
 		tree.deleteCase5(node)
@@ -363,32 +372,41 @@ func (tree *Tree) deleteCase4(node *Node) {
 }
 
 func (tree *Tree) deleteCase5(node *Node) {
+	sibling := node.sibling()
 	if node == node.parent.left &&
-		node.sibling().color == BLACK &&
+		color(sibling) == BLACK &&
-		node.sibling().left.color == RED &&
+		color(sibling.left) == RED &&
-		node.sibling().right.color == BLACK {
+		color(sibling.right) == BLACK {
-		node.sibling().color = RED
+		sibling.color = RED
-		node.sibling().left.color = BLACK
+		sibling.left.color = BLACK
-		tree.rotateRight(node.sibling())
+		tree.rotateRight(sibling)
 	} else if node == node.parent.right &&
-		node.sibling().color == BLACK &&
+		color(sibling) == BLACK &&
-		node.sibling().right.color == RED &&
+		color(sibling.right) == RED &&
-		node.sibling().left.color == BLACK {
+		color(sibling.left) == BLACK {
-		node.sibling().color = RED
+		sibling.color = RED
-		node.sibling().right.color = BLACK
+		sibling.right.color = BLACK
-		tree.rotateLeft(node.sibling())
+		tree.rotateLeft(sibling)
 	}
 	tree.deleteCase6(node)
 }
 
 func (tree *Tree) deleteCase6(node *Node) {
-	node.sibling().color = node.parent.color
+	sibling := node.sibling()
+	sibling.color = color(node.parent)
 	node.parent.color = BLACK
-	if node == node.parent.left && node.sibling().right.color == RED {
+	if node == node.parent.left && color(sibling.right) == RED {
-		node.sibling().right.color = BLACK
+		sibling.right.color = BLACK
 		tree.rotateLeft(node.parent)
-	} else if node.sibling().left.color == RED {
+	} else if color(sibling.left) == RED {
-		node.sibling().left.color = BLACK
+		sibling.left.color = BLACK
 		tree.rotateRight(node.parent)
 	}
 }
+
+func color(node *Node) Color {
+	if node == nil {
+		return BLACK
+	}
+	return node.color
+}