add Size to search interface and fix redblacktree tests

1 files changed, 9 insertions, 0 deletions

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