avl_tree: adding count variable to keep the number of nodes in tree.

src/lol/algorithm/avl_tree.h

@@ -20,12 +20,14 @@ class avl_tree
 {
 public:
     avl_tree() :
-        m_root(nullptr)
+        m_root(nullptr),
+        m_count(0)
     {
     }
 
     avl_tree(avl_tree const & other) :
-        m_root(nullptr)
+        m_root(nullptr),
+        m_count(0)
     {
         for (auto iterator : other)
             this->insert(iterator.key, iterator.value);
@@ -40,6 +42,8 @@ public:
             this->m_root = nullptr;
         }
 
+        this->m_count = other->m_count;
+
         for (auto iterator : other)
             this->insert(iterator.key, iterator.value);
     }
@@ -58,10 +62,17 @@ public:
         if (!m_root)
         {
             this->m_root = new tree_node(key, value, &this->m_root);
+            ++this->m_count;
+            return true;
+        }
+
+        if(this->m_root->insert(key, value))
+        {
+            ++this->m_count;
             return true;
         }
 
-        return this->m_root->insert(key, value);
+        return false;
     }
 
     bool erase(K const & key)
@@ -69,7 +80,13 @@ public:
         if (!m_root)
             return false;
 
-        return this->m_root->erase(key);
+        if(this->m_root->erase(key))
+        {
+            --this->m_count;
+            return true;
+        }
+
+        return false;
     }
 
     bool exists(K const & key)
@@ -80,6 +97,16 @@ public:
         return this->m_root->exists(key);
     }
 
+    void clear()
+    {
+        if (this->m_root)
+        {
+            this->m_root->cascade_delete();
+            delete this->m_root;
+            this->m_root = nullptr;
+        }
+    }
+
     bool try_get_value(K const & key, V * & value_ptr) const
     {
         if (this->m_root)
@@ -143,6 +170,11 @@ public:
         return ConstIterator(node);
     }
 
+    int get_count() const
+    {
+        return this->m_count;
+    }
+
     Iterator end()
     {
         return Iterator(nullptr);
@@ -559,6 +591,8 @@ public:
 protected:
 
     tree_node * m_root;
+
+    int m_count;
 };
 
 }

src/t/algorithm/avl_tree.cpp

@@ -88,12 +88,16 @@ lolunit_declare_fixture(AvlTreeTest)
         lolunit_assert_equal(tree.insert(30, 1), true);
         lolunit_assert_equal(tree.get_root_balance(), 0);
 
+        lolunit_assert_equal(tree.get_count(), 3);
+
         lolunit_assert_equal(tree.erase(30), true);
+        lolunit_assert_equal(tree.get_count(), 2);
         lolunit_assert_equal(tree.insert(30, 1), true);
 
         lolunit_assert_equal(tree.erase(20), true);
         lolunit_assert_equal(tree.insert(20, 1), true);
         lolunit_assert_equal(tree.exists(10), true);
+
     }
 
     lolunit_declare_test(AvlTreeExistence)
@@ -189,6 +193,8 @@ lolunit_declare_fixture(AvlTreeTest)
             lolunit_assert_equal(iterator.value == (iterator.key * 2 + iterator.key % 3), true);
             tmp = iterator.key;
         }
+
+        lolunit_assert_equal(tree.get_count(), 99);
     }
 
     lolunit_declare_test(AvlTreeTestIteratorCopy)
@@ -208,6 +214,8 @@ lolunit_declare_fixture(AvlTreeTest)
             lolunit_assert_equal(iterator.value == (iterator.key * 2 + iterator.key % 3), true);
             tmp = iterator.key;
         }
+
+        lolunit_assert_equal(other.get_count(), 99);
     }
 };