From 7c76ce4bce05366e8d43ff1595b03a98d73223b8 Mon Sep 17 00:00:00 2001
From: Guillaume Bittoun <guillaume.bittoun@gmail.com>
Date: Tue, 14 Oct 2014 10:58:14 +0000
Subject: [PATCH] avl_tree: adding mandatory get method

---
 src/lol/algorithm/avl_tree.h | 24 ++++++++++++++++++
 src/t/algorithm/avl_tree.cpp | 49 +++++++++++++++++++++++++++++++++++-
 2 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/src/lol/algorithm/avl_tree.h b/src/lol/algorithm/avl_tree.h
index f0769df4..47bd17aa 100644
--- a/src/lol/algorithm/avl_tree.h
+++ b/src/lol/algorithm/avl_tree.h
@@ -52,6 +52,14 @@ public:
         return this->m_root->exists(key);
     }
 
+    bool try_get_value(K const & key, V * & value_ptr)
+    {
+        if (this->m_root)
+            return this->m_root->try_get_value(key, value_ptr);
+
+        return false;
+    }
+
     virtual ~avl_tree()
     {
         if (this->m_root)
@@ -122,6 +130,22 @@ protected:
             return false;
         }
 
+        bool try_get_value(K const & key, V * & value_ptr)
+        {
+            int i = -1 + (key < this->m_key) + 2 * (this->m_key < key);
+
+            if (i < 0)
+            {
+                value_ptr = &this->m_value;
+                return true;
+            }
+
+            if (this->m_child[i])
+                return this->m_child[i]->try_get_value(key, value_ptr);
+
+            return false;
+        }
+
         tree_node * detach_deepest(int i, tree_node * & parent_slot)
         {
             tree_node * ret = nullptr;
diff --git a/src/t/algorithm/avl_tree.cpp b/src/t/algorithm/avl_tree.cpp
index 1f510eab..90713f26 100644
--- a/src/t/algorithm/avl_tree.cpp
+++ b/src/t/algorithm/avl_tree.cpp
@@ -71,7 +71,7 @@ lolunit_declare_fixture(AvlTreeTest)
         lolunit_assert_equal(tree.insert(11, 1), true);
         lolunit_assert_equal(tree.get_root_balance(), 1);
 
-        lolunit_assert_equal(tree.insert(13, 1), true);
+        lolunit_assert_equal(tree.insert(13, 2), true);
         lolunit_assert_equal(tree.get_root_balance(), 1);
     }
 
@@ -125,6 +125,53 @@ lolunit_declare_fixture(AvlTreeTest)
         lolunit_assert_equal(tree.exists(40), true);
         lolunit_assert_equal(tree.exists(50), true);
     }
+
+
+    lolunit_declare_test(AvlTreeGetValue)
+    {
+        test_tree tree;
+
+        lolunit_assert_equal(tree.insert(10, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), 0);
+
+        lolunit_assert_equal(tree.insert(20, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), 1);
+
+        lolunit_assert_equal(tree.insert(30, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), 0);
+
+        lolunit_assert_equal(tree.insert(00, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), -1);
+
+        lolunit_assert_equal(tree.insert(-10, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), -1);
+
+        lolunit_assert_equal(tree.insert(-20, 1), true);
+        lolunit_assert_equal(tree.get_root_balance(), 0);
+
+        lolunit_assert_equal(tree.insert(-20, 1), false);
+        lolunit_assert_equal(tree.get_root_balance(), 0);
+
+        lolunit_assert_equal(tree.insert(11, 2), true);
+        lolunit_assert_equal(tree.get_root_balance(), 1);
+
+        lolunit_assert_equal(tree.insert(13, 3), true);
+        lolunit_assert_equal(tree.get_root_balance(), 1);
+
+        int * value_ptr = nullptr;
+
+        lolunit_assert_equal(tree.try_get_value(-10, value_ptr), true);
+        lolunit_assert_equal(*value_ptr, 1);
+
+        lolunit_assert_equal(tree.try_get_value(11, value_ptr), true);
+        lolunit_assert_equal(*value_ptr, 2);
+
+        lolunit_assert_equal(tree.try_get_value(13, value_ptr), true);
+        lolunit_assert_equal(*value_ptr, 3);
+
+        lolunit_assert_equal(tree.try_get_value(67, value_ptr), false);
+        lolunit_assert_equal(*value_ptr, 3);
+    }
 };
 
 }