initial del support for redblack tree

2 files changed, 75 insertions, 29 deletions

diff --git a/search/redblackbst.go b/search/redblackbst.go
index c1d0600..6644aea 100644
--- a/search/redblackbst.go
+++ b/search/redblackbst.go
@@ -1,6 +1,8 @@
 package search
 
-import "fmt"
+import (
+	"fmt"
+)
 
 type linkColor bool
 
@@ -10,12 +12,14 @@ const (
 )
 
 type rbNode struct {
-	key   int
-	val   int
-	color linkColor
-	size  int
-	left  *rbNode
-	right *rbNode
+	key      int
+	val      int
+	color    linkColor
+	capacity int
+	left     *rbNode
+	right    *rbNode
+	// Just mark a node as deleted if deleted. Not fully implemented in lecture.
+	deleted bool
 }
 
 func (n *rbNode) String() string {
@@ -23,14 +27,19 @@ func (n *rbNode) String() string {
 		if n == nil {
 			return ""
 		}
-		return n.String()
+		return fmt.Sprintf("\n%s", n.String())
 	}
 
-	return fmt.Sprintf("rbNode{%d:%d,%v,%d,%s,%s}",
+	color := "red"
+	if n.color == black {
+		color = "black"
+	}
+	return fmt.Sprintf("rbNode{%v;%d:%d,%s,%d,%s,%s}",
+		n.deleted,
 		n.key,
 		n.val,
-		n.color,
-		n.size,
+		color,
+		n.capacity,
 		recurse(n.left),
 		recurse(n.right),
 	)
@@ -43,15 +52,16 @@ func (n *rbNode) isRed() bool {
 	return n.color == red
 }
 
-func (n *rbNode) Size() int {
+func (n *rbNode) Capacity() int {
 	if n == nil {
 		return 0
 	}
-	return n.size
+	return n.capacity
 }
 
 type RedBlackBST struct {
 	root *rbNode
+	size int
 }
 
 func NewRedBlackBST() *RedBlackBST {
@@ -60,21 +70,25 @@ func NewRedBlackBST() *RedBlackBST {
 
 func (t *RedBlackBST) String() string {
 	if t.Empty() {
-		return "RedBlackBST{}"
+		return fmt.Sprintf("RedBlackBST{%d:%d}", t.Size(), t.Capacity())
 	}
 
-	return fmt.Sprintf("RedBlackBST{%s}", t.root)
+	return fmt.Sprintf("RedBlackBST{%d:%d;%s}", t.Size(), t.Capacity(), t.root)
+}
+
+func (t *RedBlackBST) Size() int {
+	return t.size
 }
 
 func (t *RedBlackBST) Empty() bool {
-	return t.root == nil
+	return t.Size() == 0
 }
 
-func (t *RedBlackBST) Size() int {
-	if t.Empty() {
+func (t *RedBlackBST) Capacity() int {
+	if t.root == nil {
 		return 0
 	}
-	return t.root.size
+	return t.root.capacity
 }
 
 func (t *RedBlackBST) Put(key, val int) {
@@ -84,7 +98,8 @@ func (t *RedBlackBST) Put(key, val int) {
 
 func (t *RedBlackBST) put(n *rbNode, key, val int) *rbNode {
 	if n == nil {
-		return &rbNode{key, val, red, 1, nil, nil}
+		t.size++
+		return &rbNode{key, val, red, 1, nil, nil, false}
 	}
 
 	switch {
@@ -93,6 +108,10 @@ func (t *RedBlackBST) put(n *rbNode, key, val int) *rbNode {
 	case key > n.key:
 		n.right = t.put(n.right, key, val)
 	default:
+		if n.deleted {
+			n.deleted = false
+		}
+		t.size++
 		n.val = val
 	}
 
@@ -105,7 +124,7 @@ func (t *RedBlackBST) put(n *rbNode, key, val int) *rbNode {
 		t.flipColors(n)
 	}
 
-	n.size = 1 + n.left.Size() + n.right.Size()
+	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()
 	return n
 }
 
@@ -124,12 +143,37 @@ func (t *RedBlackBST) get(n *rbNode, key int) (int, error) {
 	case key > n.key:
 		return t.get(n.right, key)
 	default:
+		if n.deleted {
+			return 0, NotFound
+		}
 		return n.val, nil
 	}
 }
 
 func (t *RedBlackBST) Del(key int) (int, error) {
-	panic("Not yet implemented")
+	return t.del(t.root, key)
+}
+
+func (t *RedBlackBST) del(n *rbNode, key int) (int, error) {
+	if n == nil {
+		return 0, NotFound
+	}
+
+	switch {
+	case key < n.key:
+		return t.del(n.left, key)
+	case key > n.key:
+		return t.del(n.right, key)
+	default:
+		if n.deleted {
+			return 0, NotFound
+		}
+		t.size--
+		n.deleted = true
+		val := n.val
+		n.val = -1
+		return val, nil
+	}
 }
 
 func (t *RedBlackBST) rotateLeft(n *rbNode) *rbNode {
@@ -140,8 +184,8 @@ func (t *RedBlackBST) rotateLeft(n *rbNode) *rbNode {
 	x.color = n.color
 	n.color = red
 
-	x.size = n.size
-	n.size = 1 + n.left.Size() + n.right.Size()
+	x.capacity = n.capacity
+	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()
 
 	return x
 }
@@ -154,8 +198,8 @@ func (t *RedBlackBST) rotateRight(n *rbNode) *rbNode {
 	x.color = n.color
 	n.color = red
 
-	x.size = n.size
-	n.size = 1 + n.left.Size() + n.right.Size()
+	x.capacity = n.capacity
+	n.capacity = 1 + n.left.Capacity() + n.right.Capacity()
 
 	return x
 }
diff --git a/search/search_test.go b/search/search_test.go
index bf1d1b5..09ed906 100644
--- a/search/search_test.go
+++ b/search/search_test.go
@@ -10,7 +10,9 @@ import (
 
 const factor int = 10
 const minLength int = 1
-const maxLength int = 10000
+const maxLength int = 10
+
+//const maxLength int = 10000
 
 // Store results here to avoid compiler optimizations
 var benchResult int
@@ -39,10 +41,10 @@ func test(s Put, l int, t *testing.T) {
 		case del:
 			defer delete(mapping, key)
 			val, err = s.Del(key)
-			//t.Log("Del", key, val, err)
+			t.Log("Del", key, val, err)
 		default:
 			val, err = s.Get(key)
-			//t.Log("Get", key, val, err)
+			t.Log("Get", key, val, err)
 		}
 
 		if mVal, ok := mapping[key]; ok {