initial refactor of the red-black tree

1 files changed, 74 insertions, 116 deletions

diff --git a/search/redblackbst.go b/search/redblackbst.go
index 167f567..c1d0600 100644
--- a/search/redblackbst.go
+++ b/search/redblackbst.go
@@ -12,9 +12,10 @@ const (
 type rbNode struct {
 	key   int
 	val   int
+	color linkColor
+	size  int
 	left  *rbNode
 	right *rbNode
-	color linkColor
 }
 
 func (n *rbNode) String() string {
@@ -25,15 +26,30 @@ func (n *rbNode) String() string {
 		return n.String()
 	}
 
-	return fmt.Sprintf("rbNode{%d:%d,%v,%s,%s}",
+	return fmt.Sprintf("rbNode{%d:%d,%v,%d,%s,%s}",
 		n.key,
 		n.val,
 		n.color,
+		n.size,
 		recurse(n.left),
 		recurse(n.right),
 	)
 }
 
+func (n *rbNode) isRed() bool {
+	if n == nil {
+		return false
+	}
+	return n.color == red
+}
+
+func (n *rbNode) Size() int {
+	if n == nil {
+		return 0
+	}
+	return n.size
+}
+
 type RedBlackBST struct {
 	root *rbNode
 }
@@ -54,156 +70,98 @@ func (t *RedBlackBST) Empty() bool {
 	return t.root == nil
 }
 
-func (t *RedBlackBST) Put(key, val int) {
+func (t *RedBlackBST) Size() int {
 	if t.Empty() {
-		t.root = &rbNode{key, val, nil, nil, black}
-		return
-	}
-	parent, ptr, _, err := t.search(nil, &t.root, key)
-	switch err {
-	case nil:
-		// key already in the tree
-		return
-	case NotFound:
-		*ptr = &rbNode{key, val, nil, nil, red}
-		return
-	default:
-		panic(err)
+		return 0
 	}
-
-	t.balance(parent)
+	return t.root.size
 }
 
-func (t *RedBlackBST) balance(n *rbNode) {
-	switch {
-	case n == nil, n.right == nil:
-		return
-	case n.right.color == black:
-		return
-	case n.left.color == black:
-		// Left is black and right is red
-	case n.left.color == red:
-		// Left and right are red
-	}
-}
-
-func (t *RedBlackBST) Get(key int) (int, error) {
-	_, _, n, err := t.search(nil, &t.root, key)
-	if err != nil {
-		return 0, err
-	}
-	return n.val, nil
+func (t *RedBlackBST) Put(key, val int) {
+	t.root = t.put(t.root, key, val)
+	t.root.color = black
 }
 
-func (t *RedBlackBST) Del(key int) (int, error) {
-	_, ptr, n, err := t.search(nil, &t.root, key)
-	if err != nil {
-		return 0, err
+func (t *RedBlackBST) put(n *rbNode, key, val int) *rbNode {
+	if n == nil {
+		return &rbNode{key, val, red, 1, nil, nil}
 	}
 
-	// Case 1: n is leaf rbNode
-	// Case 2: n has one child
-	// Case 3: n has two childs
-
 	switch {
-	case n.left == nil:
-		if n.right == nil {
-			// I am a leaf rbNode
-			*ptr = nil
-			return n.val, nil
-		}
-		// I have a right child
-		*ptr = n.right
-		return n.val, nil
-
-	case n.right == nil:
-		// I have a left child
-		*ptr = n.left
-		return n.val, nil
+	case key < n.key:
+		n.left = t.put(n.left, key, val)
+	case key > n.key:
+		n.right = t.put(n.right, key, val)
 	default:
-		// I have two children!
+		n.val = val
+	}
 
-		o, err := t.deleteMin(&n.right)
-		if err != nil {
-			return 0, err
-		}
+	switch {
+	case n.right.isRed() && !n.left.isRed():
+		n = t.rotateLeft(n)
+	case n.left != nil && n.left.isRed() && n.left.left.isRed():
+		n = t.rotateRight(n)
+	case n.left.isRed() && n.right.isRed():
+		t.flipColors(n)
+	}
 
-		o.left = n.left
-		o.right = n.right
-		*ptr = o
+	n.size = 1 + n.left.Size() + n.right.Size()
+	return n
+}
 
-		return n.val, nil
-	}
+func (t *RedBlackBST) Get(key int) (int, error) {
+	return t.get(t.root, key)
 }
 
-func (t *RedBlackBST) search(parent *rbNode, ptr **rbNode, key int) (*rbNode, **rbNode, *rbNode, error) {
-	n := *ptr
+func (t *RedBlackBST) get(n *rbNode, key int) (int, error) {
 	if n == nil {
-		return parent, ptr, nil, NotFound
+		return 0, NotFound
 	}
 
 	switch {
 	case key < n.key:
-		return t.search(n, &n.left, key)
-	case n.key < key:
-		return t.search(n, &n.right, key)
+		return t.get(n.left, key)
+	case key > n.key:
+		return t.get(n.right, key)
 	default:
-		return parent, ptr, n, nil
+		return n.val, nil
 	}
 }
 
-func (t *RedBlackBST) deleteMin(ptr **rbNode) (*rbNode, error) {
-	ptr, n, err := t.min(ptr)
-	if err != nil {
-		return nil, err
-	}
-
-	*ptr = n.right
-	n.right = nil
-	return n, nil
+func (t *RedBlackBST) Del(key int) (int, error) {
+	panic("Not yet implemented")
 }
 
-func (t *RedBlackBST) min(ptr **rbNode) (**rbNode, *rbNode, error) {
-	n := *ptr
-	if n == nil {
-		return nil, nil, NotFound
-	}
+func (t *RedBlackBST) rotateLeft(n *rbNode) *rbNode {
+	x := n.right
+	n.right = x.left
+	x.left = n
 
-	for {
-		if n.left == nil {
-			return ptr, n, nil
-		}
-		ptr = &n.left
-		n = *ptr
-	}
-}
-
-func (t *RedBlackBST) rotateLeft(ptr **rbNode) {
-	x := *ptr
-	y := x.right
+	x.color = n.color
+	n.color = red
 
-	x.right = y.left
-	x.left = y
+	x.size = n.size
+	n.size = 1 + n.left.Size() + n.right.Size()
 
-	*ptr = y
-	x.color = red
-	y.color = black
+	return x
 }
 
-func (t *RedBlackBST) rotateRight(ptr **rbNode) {
-	y := *ptr
-	x := y.right
+func (t *RedBlackBST) rotateRight(n *rbNode) *rbNode {
+	x := n.left
+	n.left = x.right
+	x.right = n
+
+	x.color = n.color
+	n.color = red
 
-	y.left = x.right
-	x.right = y
+	x.size = n.size
+	n.size = 1 + n.left.Size() + n.right.Size()
 
-	*ptr = x
-	x.color = black
-	y.color = red
+	return x
 }
 
 func (t *RedBlackBST) flipColors(n *rbNode) {
+	n.color = red
 	n.left.color = black
 	n.right.color = black
-	n.color = red
 }