diff --git a/src/lol/algorithm/avl_tree.h b/src/lol/algorithm/avl_tree.h index 48a22b96..c383d227 100644 --- a/src/lol/algorithm/avl_tree.h +++ b/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; diff --git a/src/t/algorithm/avl_tree.cpp b/src/t/algorithm/avl_tree.cpp index 84904704..ea55b447 100644 --- a/src/t/algorithm/avl_tree.cpp +++ b/src/t/algorithm/avl_tree.cpp @@ -12,6 +12,8 @@ #include +#include + #include namespace lol