avl_tree: phewww ! feels better… To be continued, memory still needs to be freed and some features are missing

src/lol/algorithm/avl_tree.h

@@ -28,58 +28,28 @@ public:
     bool insert(K const & key, V const & value)
     {
         if (!m_root)
-            this->m_root = new tree_node(key, value);
-        else if (this->m_root->exists(key))
-        {
-            this->m_root->insert_or_update(key, value);
-            return false;
-        }
-        else
         {
-            this->m_root->insert_or_update(key, value);
-            this->rebalance_if_needed(key);
+            this->m_root = new tree_node(key, value);
+            return true;
         }
 
-        return true;
+        return this->m_root->insert(key, value, this->m_root);
     }
 
     bool erase(K const & key)
     {
-        if (this->m_root)
-        {
-            tree_node * parent = this->m_root->get_parent(key);
-
-            if (parent)
-            {
-                parent->erase_child(key);
-                this->rebalance_if_needed(key);
-            }
-            else if (this->m_root->key_equals(key))
-            {
-                this->m_root = this->m_root->erase_self();
-                this->rebalance_if_needed(key);
-            }
-
-            return true;
-        }
+        if (!m_root)
+            return false;
 
-        return false;
+        return this->m_root->erase(key, this->m_root);
     }
 
-    void rebalance_if_needed(K const & key)
+    bool exists(K const & key)
     {
-        this->m_root->update_balance(key);
-        tree_node * node = this->m_root->get_unbalanced_parent(key);
+        if (!m_root)
+            return false;
 
-        if (node)
-        {
-            node->rebalance_children(key);
-            this->m_root->update_balance(node);
-        }
-        else if (this->m_root->get_balance() == -2)
-            this->m_root = this->m_root->rotate(tree_node::CW);
-        else if (this->m_root->get_balance() == 2)
-            this->m_root = this->m_root->rotate(tree_node::CCW);
+        return this->m_root->exists(key);
     }
 
 protected:
@@ -95,241 +65,108 @@ protected:
             m_stairs[0] = m_stairs[1] = 0;
         }
 
-        bool exists(K key)
-        {
-            int i = -1 + (key < this->m_key) + 2 * (this->m_key < key);
-
-            if (i < 0)
-                return true;
-
-            if (this->m_child[i])
-                return this->m_child[i]->exists(key);
-
-            return false;
-        }
-
-        void insert_or_update(K const & key, V const & value)
+        /* Insert a value in tree and return true or update an existing value for
+         * the existing key and return false
+         */
+        bool insert(K const & key, V const & value, tree_node * & parent_slot)
         {
             int i = -1 + (key < this->m_key) + 2 * (this->m_key < key);
 
             if (i < 0)
             {
                 this->m_value = value;
+                return false;
             }
+
+            bool created = false;
+
+            if (this->m_child[i])
+                created = this->m_child[i]->insert(key, value, this->m_child[i]);
             else
             {
-                if (this->m_child[i])
-                    this->m_child[i]->insert_or_update(key, value);
-                else
-                    this->m_child[i] = new tree_node(key, value);
+                this->m_child[i] = new tree_node(key, value);
+                created = true;
             }
-        }
 
-        void update_balance(tree_node * node)
-        {
-            this->update_balance(node->m_key);
-        }
+            if (created)
+                this->rebalance_if_needed(parent_slot);
 
-        void update_balance(K const & key)
-        {
-            if (key < this->m_key && this->m_child[0])
-                this->m_child[0]->update_balance(key);
-            if (this->m_key < key && this->m_child[1])
-                this->m_child[1]->update_balance(key);
-
-            this->compute_balance();
+            return created;
         }
 
-        /* Retrieve the parent of the deeper unbalanced node after key insertion.
-         * Do not call "get_unbalanced_parent" if key is not already present in the tree. */
-        tree_node * get_unbalanced_parent(K const & key)
+        bool erase(K const & key, tree_node * & parent_slot)
         {
-            tree_node * parent = nullptr;
+            int i = -1 + (key < this->m_key) + 2 * (this->m_key < key);
 
-            if (key < this->m_key && this->m_child[0])
+            if (i < 0)
             {
-                parent = this->m_child[0]->get_unbalanced_parent(key);
-
-                if (parent)
-                    return parent;
-                else if (abs(this->m_child[0]->get_balance()) == 2)
-                    return this;
+                this->erase_self(parent_slot);
+                delete this; // FUCK YEAH !!
+                return true;
             }
-            if (this->m_key < key && this->m_child[1])
+            else if(this->m_child[i]->erase(key, this->m_child[i]))
             {
-                parent = this->m_child[1]->get_unbalanced_parent(key);
-
-                if (parent)
-                    return parent;
-                else if (abs(this->m_child[1]->get_balance()) == 2)
-                    return this;
+                this->rebalance_if_needed(parent_slot);
+                return true;
             }
 
-            return nullptr;
+            return false;
         }
 
-        void rebalance_children(K const & key)
+        tree_node * detach_deepest(int i, tree_node * & parent_slot)
         {
-            if (key < this->m_key)
-            {
-                if (this->m_child[0]->get_balance() == 2)
-                    this->m_child[0] = this->m_child[0]->rotate(CCW);
-                else if (this->m_child[0]->get_balance() == -2)
-                    this->m_child[0] = this->m_child[0]->rotate(CW);
-            }
-            else if (this->m_key < key)
-            {
-                if (this->m_child[1]->get_balance() == 2)
-                    this->m_child[1] = this->m_child[1]->rotate(CCW);
-                else if (this->m_child[1]->get_balance() == -2)
-                    this->m_child[1] = this->m_child[1]->rotate(CW);
-            }
-            else
-                ASSERT(false) // Do not rebalance the "this" node here
-        }
-
-        enum Rotation { CW = 0, CCW = 1 };
+            tree_node * ret = nullptr;
 
-        tree_node * rotate(Rotation rotation)
-        {
-            if (rotation == CW)
+            if (this->m_child[i])
+                ret = this->m_child[i]->detach_deepest(i, this->m_child[i]);
+            else
             {
-                tree_node * newhead = this->m_child[0];
-                tree_node * tmp = this->m_child[0]->m_child[1];
-
-                this->m_child[0]->m_child[1] = this;
-                this->m_child[0] = tmp;
-
-                this->compute_balance();
-                newhead->compute_balance();
-
-                return newhead;
+                parent_slot = this->m_child[(i ? 0 : 1)];
+                this->m_child[(i ? 0 : 1)] = nullptr;
+                ret = this;
             }
-            else // rotation == CCW
-            {
-                tree_node * newhead = this->m_child[1];
-                tree_node * tmp = this->m_child[1]->m_child[0];
-
-                this->m_child[1]->m_child[0] = this;
-                this->m_child[1] = tmp;
 
-                this->compute_balance();
-                newhead->compute_balance();
+            this->rebalance_if_needed(parent_slot);
 
-                return newhead;
-            }
-
-            return 0;
+            return ret;
         }
 
-        void compute_balance()
+        void update_balance()
         {
             this->m_stairs[0] = this->m_child[0] ? lol::max(this->m_child[0]->m_stairs[0], this->m_child[0]->m_stairs[1]) + 1 : 0;
             this->m_stairs[1] = this->m_child[1] ? lol::max(this->m_child[1]->m_stairs[0], this->m_child[1]->m_stairs[1]) + 1 : 0;
         }
 
-        bool key_equals(K const & key)
-        {
-            return !(key < this->m_key) && !(this->m_key < key);
-        }
-
-        /* Retrieve the parent of a key.
-         * Do not call "get_parent" if you’re not sure the key or node is present. */
-        tree_node * get_parent(tree_node * node)
-        {
-            return this->get_parent(node->m_key);
-        }
-
-        /* Retrieve the parent of an inserted key.
-         * Do not call "get_parent" if key is not already present in the tree. */
-        tree_node * get_parent(K const & key)
-        {
-            if (key < this->m_key)
-            {
-                if (this->m_child[0]->key_equals(key))
-                    return this;
-                else
-                    return this->m_child[0]->get_parent(key);
-            }
-            else if (this->m_key < key)
-            {
-                if (this->m_child[1]->key_equals(key))
-                    return this;
-                else
-                    return this->m_child[1]->get_parent(key);
-            }
-
-            return nullptr;
-        }
-
-        void erase_child(K const & key)
+        void rebalance_if_needed(tree_node * & parent_slot)
         {
-            tree_node * erase_me = nullptr;
-
-            if (key < this->m_key)
-            {
-                erase_me = this->m_child[0];
-                this->m_child[0] = erase_me->erase_self();
-            }
-            else if (this->m_key < key)
-            {
-                erase_me = this->m_child[1];
-                this->m_child[1] = erase_me->erase_self();
-            }
-            else
-                ASSERT(false) // Do not erase the "this" node here
+            this->update_balance();
 
-            delete erase_me;
-        }
+            int i = (this->get_balance() ==  2);
+            int j = (this->get_balance() == -2);
 
-        tree_node * erase_self()
-        {
-            tree_node * replacement = nullptr;
-
-            if (this->get_balance() == -1)
-            {
-                replacement = this->get_deeper(0);
-                if (replacement)
-                    this->get_parent(replacement)->m_child[1] = replacement->m_child[0];
-            }
-            else // this->get_balance() >= 0
+            if (i || j)
             {
-                replacement = this->get_deeper(1);
-                if (replacement)
-                    this->get_parent(replacement)->m_child[0] = replacement->m_child[1];
-            }
+                tree_node * swap = this->m_child[i];
+                this->m_child[i] = swap->m_child[j];
+                swap->m_child[j] = this;
+                parent_slot = swap;
 
-            if (replacement)
-            {
-                replacement->m_child[0] = this->m_child[0];
-                replacement->m_child[1] = this->m_child[1];
+                this->update_balance();
+                swap->update_balance();
             }
-
-            return replacement;
         }
 
-        void replace(tree_node * from, tree_node * to)
+        bool exists(K const & key)
         {
-            from->m_child[0] = to->m_child[0];
-            from->m_child[1] = to->m_child[1];
-
-            to->m_child[0] = nullptr;
-            to->m_child[1] = nullptr;
-        }
+            int i = -1 + (key < this->m_key) + 2 * (this->m_key < key);
 
-        tree_node * get_deeper(int index)
-        {
-            tree_node * previous = this->m_child[index];
+            if (i < 0)
+                return true;
 
-            if (previous)
-            {
-                while (previous->m_child[1 - index])
-                {
-                    previous = previous->m_child[1 - index];
-                }
-            }
+            if (this->m_child[i])
+                return this->m_child[i]->exists(key);
 
-            return previous;
+            return false;
         }
 
         int get_balance()
@@ -344,6 +181,18 @@ protected:
 
     protected:
 
+        void erase_self(tree_node * & parent_slot)
+        {
+            int i = (this->get_balance() == -1);
+
+            tree_node * replacement = nullptr;
+
+            if (i || this->m_child[1])
+                replacement = this->m_child[1 - i]->detach_deepest(i, this->m_child[1 - i]);
+
+            parent_slot = replacement;
+        }
+
         K m_key;
         V m_value;
 

src/t/algorithm/avl_tree.cpp

@@ -12,6 +12,8 @@
 
 #include <lol/engine-internal.h>
 
+#include <lol/algorithm/avl_tree.h>
+
 #include <lolunit.h>
 
 namespace lol