add Size to search interface and fix redblacktree tests

4 files changed, 19 insertions, 18 deletions

diff --git a/search/bst.go b/search/bst.go
index b1d12bc..a1d2fe0 100644
--- a/search/bst.go
+++ b/search/bst.go
@@ -26,6 +26,7 @@ func (n *node) String() string {
 
 type BST struct {
 	root *node
+	size int
 }
 
 func NewBST() *BST {
@@ -44,9 +45,14 @@ func (t *BST) Empty() bool {
 	return t.root == nil
 }
 
+func (t *BST) Size() int {
+	return t.size
+}
+
 func (t *BST) Put(key, val int) {
 	if t.Empty() {
 		t.root = &node{key, val, nil, nil}
+		t.size++
 		return
 	}
 	ptr, _, err := t.search(&t.root, key)
@@ -56,6 +62,7 @@ func (t *BST) Put(key, val int) {
 		return
 	case NotFound:
 		*ptr = &node{key, val, nil, nil}
+		t.size++
 		return
 	default:
 		panic(err)
@@ -75,6 +82,7 @@ func (t *BST) Del(key int) (int, error) {
 	if err != nil {
 		return 0, err
 	}
+	t.size--
 
 	// Case 1: n is leaf node
 	// Case 2: n has one child
@@ -152,20 +160,3 @@ func (t *BST) min(ptr **node) (**node, *node, error) {
 		n = *ptr
 	}
 }
-
-/*
-func (t *BST) max(ptr **node) (**node, *node, error) {
-	n := *ptr
-	if n == nil {
-		return nil, nil, NotFound
-	}
-
-	for {
-		if n.right == nil {
-			return ptr, n, nil
-		}
-		ptr = &n.right
-		n = *ptr
-	}
-}
-*/
diff --git a/search/elementary.go b/search/elementary.go
index b522cd0..ed00159 100644
--- a/search/elementary.go
+++ b/search/elementary.go
@@ -8,6 +8,7 @@ type ElementaryElem struct {
 
 type Elementary struct {
 	root *ElementaryElem
+	size int
 }
 
 func NewElementary() *Elementary {
@@ -18,9 +19,14 @@ func (s *Elementary) Empty() bool {
 	return s.root == nil
 }
 
+func (s *Elementary) Size() int {
+	return s.size
+}
+
 func (s *Elementary) Put(key int, val int) {
 	if s.Empty() {
 		s.root = &ElementaryElem{key, val, nil}
+		s.size++
 		return
 	}
 
@@ -33,6 +39,7 @@ func (s *Elementary) Put(key int, val int) {
 		}
 		if elem.next == nil {
 			elem.next = &ElementaryElem{key, val, nil}
+			s.size++
 			return
 		}
 		elem = elem.next
@@ -61,6 +68,7 @@ func (s *Elementary) Del(key int) (int, error) {
 		defer func() {
 			s.root = s.root.next
 		}()
+		s.size--
 		return s.root.val, nil
 	}
 
@@ -70,6 +78,7 @@ func (s *Elementary) Del(key int) (int, error) {
 			defer func() {
 				elem.next = elem.next.next
 			}()
+			s.size--
 			return elem.next.val, nil
 		}
 		elem = elem.next
diff --git a/search/redblackbst.go b/search/redblackbst.go
index 6644aea..856d562 100644
--- a/search/redblackbst.go
+++ b/search/redblackbst.go
@@ -110,8 +110,8 @@ func (t *RedBlackBST) put(n *rbNode, key, val int) *rbNode {
 	default:
 		if n.deleted {
 			n.deleted = false
+			t.size++
 		}
-		t.size++
 		n.val = val
 	}
 
diff --git a/search/set.go b/search/set.go
index 17024ea..407ca44 100644
--- a/search/set.go
+++ b/search/set.go
@@ -9,6 +9,7 @@ var (
 
 type Put interface {
 	Empty() bool
+	Size() int
 	Put(key int, val int)
 	Get(key int) (int, error)
 	Del(key int) (int, error)