From 3ba3ce1b0ef8864c740d37a28740faf57bfdbef1 Mon Sep 17 00:00:00 2001
From: Emir Pasic <emirpasic@hotmail.com>
Date: Thu, 5 Mar 2015 05:20:10 +0100
Subject: [PATCH] - adding Size() method to the redblacktree

---
 trees/redblacktree/redblacktree.go | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/trees/redblacktree/redblacktree.go b/trees/redblacktree/redblacktree.go
index 6aedd72..88d29f7 100644
--- a/trees/redblacktree/redblacktree.go
+++ b/trees/redblacktree/redblacktree.go
@@ -35,6 +35,7 @@ const (
 
 type Tree struct {
 	root       *Node
+	size       int
 	comparator utils.Comparator
 }
 
@@ -96,6 +97,7 @@ func (tree *Tree) Put(key interface{}, value interface{}) {
 		insertedNode.parent = node
 	}
 	tree.insertCase1(insertedNode)
+	tree.size += 1
 }
 
 // Searches the node in the tree by key and returns its value or nil if key is not found in tree.
@@ -138,6 +140,7 @@ func (tree *Tree) Remove(key interface{}) {
 			child.color = BLACK
 		}
 	}
+	tree.size -= 1
 }
 
 // Returns true if tree does not contain any nodes
@@ -145,6 +148,11 @@ func (tree *Tree) IsEmpty() bool {
 	return tree.root == nil
 }
 
+// Returns number of nodes in the tree.
+func (tree *Tree) Size() int {
+	return tree.size
+}
+
 func (tree *Tree) String() string {
 	str := "RedBlackTree\n"
 	if !tree.IsEmpty() {