diff --git a/core/map.h b/core/map.h index a37d898a9cb..f01062ebed5 100644 --- a/core/map.h +++ b/core/map.h @@ -31,6 +31,7 @@ #define MAP_H #include "set.h" + /** @author Juan Linietsky */ @@ -52,7 +53,6 @@ public: private: friend class Map; - //Color color; int color; Element *right; Element *left; @@ -61,7 +61,6 @@ public: Element *_prev; K _key; V _value; - //_Data *data; public: @@ -147,7 +146,6 @@ private: #ifdef GLOBALNIL_DISABLED memdelete_allocator(_nil); #endif - //memdelete_allocator(_root); } }; @@ -158,6 +156,7 @@ private: ERR_FAIL_COND(p_node == _data._nil && p_color == RED); p_node->color = p_color; } + inline void _rotate_left(Element *p_node) { Element *r = p_node->right; @@ -206,8 +205,9 @@ private: while (node == node->parent->right) { node = node->parent; } + if (node->parent == _data._root) - return NULL; + return NULL; // No successor, as p_node = last node return node->parent; } } @@ -225,10 +225,11 @@ private: } else { while (node == node->parent->left) { - if (node->parent == _data._root) - return NULL; node = node->parent; } + + if (node == _data._root) + return NULL; // No predecessor, as p_node = first node return node->parent; } } @@ -239,16 +240,15 @@ private: C less; while (node != _data._nil) { - if (less(p_key, node->_key)) node = node->left; else if (less(node->_key, p_key)) node = node->right; else - break; // found + return node; // found } - return (node != _data._nil) ? node : NULL; + return NULL; } Element *_find_closest(const K &p_key) const { @@ -265,24 +265,68 @@ private: else if (less(node->_key, p_key)) node = node->right; else - break; // found + return node; // found } - if (node == _data._nil) { - if (prev == NULL) - return NULL; - if (less(p_key, prev->_key)) { + if (prev == NULL) + return NULL; // tree empty - prev = prev->_prev; - } + if (less(p_key, prev->_key)) + prev = prev->_prev; - return prev; - - } else - return node; + return prev; } - Element *_insert(const K &p_key, bool &r_exists) { + void _insert_rb_fix(Element *p_new_node) { + + Element *node = p_new_node; + Element *nparent = node->parent; + Element *ngrand_parent; + + while (nparent->color == RED) { + ngrand_parent = nparent->parent; + + if (nparent == ngrand_parent->left) { + if (ngrand_parent->right->color == RED) { + _set_color(nparent, BLACK); + _set_color(ngrand_parent->right, BLACK); + _set_color(ngrand_parent, RED); + node = ngrand_parent; + nparent = node->parent; + } else { + if (node == nparent->right) { + _rotate_left(nparent); + node = nparent; + nparent = node->parent; + } + _set_color(nparent, BLACK); + _set_color(ngrand_parent, RED); + _rotate_right(ngrand_parent); + } + } else { + if (ngrand_parent->left->color == RED) { + _set_color(nparent, BLACK); + _set_color(ngrand_parent->left, BLACK); + _set_color(ngrand_parent, RED); + node = ngrand_parent; + nparent = node->parent; + } else { + if (node == nparent->left) { + _rotate_right(nparent); + node = nparent; + nparent = node->parent; + } + _set_color(nparent, BLACK); + _set_color(ngrand_parent, RED); + _rotate_left(ngrand_parent); + } + } + } + + _set_color(_data._root->left, BLACK); + } + + Element *_insert(const K &p_key, const V &p_value) { Element *new_parent = _data._root; Element *node = _data._root->left; @@ -297,27 +341,25 @@ private: else if (less(node->_key, p_key)) node = node->right; else { - r_exists = true; - return node; + node->_value = p_value; + return node; // Return existing node with new value } } Element *new_node = memnew_allocator(Element, A); - new_node->parent = new_parent; new_node->right = _data._nil; new_node->left = _data._nil; new_node->_key = p_key; + new_node->_value = p_value; //new_node->data=_data; - if (new_parent == _data._root || less(p_key, new_parent->_key)) { + if (new_parent == _data._root || less(p_key, new_parent->_key)) { new_parent->left = new_node; } else { new_parent->right = new_node; } - r_exists = false; - new_node->_next = _successor(new_node); new_node->_prev = _predecessor(new_node); if (new_node->_next) @@ -325,168 +367,113 @@ private: if (new_node->_prev) new_node->_prev->_next = new_node; - return new_node; - } - - Element *_insert_rb(const K &p_key, const V &p_value) { - - bool exists = false; - Element *new_node = _insert(p_key, exists); - - if (new_node) { - new_node->_value = p_value; - } - if (exists) - return new_node; - - Element *node = new_node; _data.size_cache++; - - while (node->parent->color == RED) { - - if (node->parent == node->parent->parent->left) { - - Element *aux = node->parent->parent->right; - - if (aux->color == RED) { - _set_color(node->parent, BLACK); - _set_color(aux, BLACK); - _set_color(node->parent->parent, RED); - node = node->parent->parent; - } else { - if (node == node->parent->right) { - node = node->parent; - _rotate_left(node); - } - _set_color(node->parent, BLACK); - _set_color(node->parent->parent, RED); - _rotate_right(node->parent->parent); - } - } else { - Element *aux = node->parent->parent->left; - - if (aux->color == RED) { - _set_color(node->parent, BLACK); - _set_color(aux, BLACK); - _set_color(node->parent->parent, RED); - node = node->parent->parent; - } else { - if (node == node->parent->left) { - node = node->parent; - _rotate_right(node); - } - _set_color(node->parent, BLACK); - _set_color(node->parent->parent, RED); - _rotate_left(node->parent->parent); - } - } - } - _set_color(_data._root->left, BLACK); + _insert_rb_fix(new_node); return new_node; } - void _erase_fix(Element *p_node) { + void _erase_fix_rb(Element *p_node) { Element *root = _data._root->left; - Element *node = p_node; + Element *node = _data._nil; + Element *sibling = p_node; + Element *parent = sibling->parent; - while ((node->color == BLACK) && (root != node)) { - if (node == node->parent->left) { - Element *aux = node->parent->right; - if (aux->color == RED) { - _set_color(aux, BLACK); - _set_color(node->parent, RED); - _rotate_left(node->parent); - aux = node->parent->right; - } - if ((aux->right->color == BLACK) && (aux->left->color == BLACK)) { - _set_color(aux, RED); - node = node->parent; + while (node != root) { // If red node found, will exit at a break + if (sibling->color == RED) { + _set_color(sibling, BLACK); + _set_color(parent, RED); + if (sibling == parent->right) { + sibling = sibling->left; + _rotate_left(parent); } else { - if (aux->right->color == BLACK) { - _set_color(aux->left, BLACK); - _set_color(aux, RED); - _rotate_right(aux); - aux = node->parent->right; + sibling = sibling->right; + _rotate_right(parent); + } + } + if ((sibling->left->color == BLACK) && (sibling->right->color == BLACK)) { + _set_color(sibling, RED); + if (parent->color == RED) { + _set_color(parent, BLACK); + break; + } else { // loop: haven't found any red nodes yet + node = parent; + parent = node->parent; + sibling = (node == parent->left) ? parent->right : parent->left; + } + } else { + if (sibling == parent->right) { + if (sibling->right->color == BLACK) { + _set_color(sibling->left, BLACK); + _set_color(sibling, RED); + _rotate_right(sibling); + sibling = sibling->parent; } - _set_color(aux, node->parent->color); - _set_color(node->parent, BLACK); - _set_color(aux->right, BLACK); - _rotate_left(node->parent); - node = root; /* this is to exit while loop */ - } - } else { /* the code below is has left and right switched from above */ - Element *aux = node->parent->left; - if (aux->color == RED) { - _set_color(aux, BLACK); - _set_color(node->parent, RED); - _rotate_right(node->parent); - aux = node->parent->left; - } - if ((aux->right->color == BLACK) && (aux->left->color == BLACK)) { - _set_color(aux, RED); - node = node->parent; + _set_color(sibling, parent->color); + _set_color(parent, BLACK); + _set_color(sibling->right, BLACK); + _rotate_left(parent); + break; } else { - if (aux->left->color == BLACK) { - _set_color(aux->right, BLACK); - _set_color(aux, RED); - _rotate_left(aux); - aux = node->parent->left; + if (sibling->left->color == BLACK) { + _set_color(sibling->right, BLACK); + _set_color(sibling, RED); + _rotate_left(sibling); + sibling = sibling->parent; } - _set_color(aux, node->parent->color); - _set_color(node->parent, BLACK); - _set_color(aux->left, BLACK); - _rotate_right(node->parent); - node = root; + + _set_color(sibling, parent->color); + _set_color(parent, BLACK); + _set_color(sibling->left, BLACK); + _rotate_right(parent); + break; } } } - _set_color(node, BLACK); - ERR_FAIL_COND(_data._nil->color != BLACK); } void _erase(Element *p_node) { - Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : _successor(p_node); - if (!rp) - rp = _data._nil; + Element *rp = ((p_node->left == _data._nil) || (p_node->right == _data._nil)) ? p_node : p_node->_next; Element *node = (rp->left == _data._nil) ? rp->right : rp->left; node->parent = rp->parent; - if (_data._root == node->parent) { - _data._root->left = node; + Element *sibling; + if (rp == rp->parent->left) { + rp->parent->left = node; + sibling = rp->parent->right; } else { - if (rp == rp->parent->left) { - rp->parent->left = node; - } else { - rp->parent->right = node; - } + rp->parent->right = node; + sibling = rp->parent->left; + } + + if (node->color == RED) { + node->parent = rp->parent; + _set_color(node, BLACK); + } else if (rp->color == BLACK && rp->parent != _data._root) { + _erase_fix_rb(sibling); } if (rp != p_node) { ERR_FAIL_COND(rp == _data._nil); - if (rp->color == BLACK) - _erase_fix(node); - rp->left = p_node->left; rp->right = p_node->right; rp->parent = p_node->parent; rp->color = p_node->color; - p_node->left->parent = rp; - p_node->right->parent = rp; + if (p_node->left != _data._nil) + p_node->left->parent = rp; + if (p_node->right != _data._nil) + p_node->right->parent = rp; if (p_node == p_node->parent->left) { p_node->parent->left = rp; } else { p_node->parent->right = rp; } - } else { - if (p_node->color == BLACK) - _erase_fix(node); } if (p_node->_next) @@ -501,11 +488,12 @@ private: void _calculate_depth(Element *p_element, int &max_d, int d) const { - if (p_element == _data._nil) { + if (p_element == _data._nil) return; - } + _calculate_depth(p_element->left, max_d, d + 1); _calculate_depth(p_element->right, max_d, d + 1); + if (d > max_d) max_d = d; } @@ -544,6 +532,7 @@ public: if (!_data._root) return NULL; + Element *res = _find(p_key); return res; } @@ -552,6 +541,7 @@ public: if (!_data._root) return NULL; + const Element *res = _find_closest(p_key); return res; } @@ -560,21 +550,28 @@ public: if (!_data._root) return NULL; + Element *res = _find_closest(p_key); return res; } + bool has(const K &p_key) const { + + return find(p_key) != NULL; + } + Element *insert(const K &p_key, const V &p_value) { if (!_data._root) _data._create_root(); - return _insert_rb(p_key, p_value); + return _insert(p_key, p_value); } void erase(Element *p_element) { - if (!_data._root) + if (!_data._root || !p_element) return; + _erase(p_element); if (_data.size_cache == 0 && _data._root) _data._free_root(); @@ -584,20 +581,17 @@ public: if (!_data._root) return false; + Element *e = find(p_key); if (!e) return false; + _erase(e); + if (_data.size_cache == 0 && _data._root) + _data._free_root(); return true; } - bool has(const K &p_key) const { - - if (!_data._root) - return false; - return find(p_key) != NULL; - } - const V &operator[](const K &p_key) const { CRASH_COND(!_data._root); @@ -605,6 +599,7 @@ public: CRASH_COND(!e); return e->_value; } + V &operator[](const K &p_key) { if (!_data._root) @@ -614,7 +609,6 @@ public: if (!e) e = insert(p_key, V()); - CRASH_COND(!e); return e->_value; } @@ -637,6 +631,7 @@ public: if (!_data._root) return NULL; + Element *e = _data._root->left; if (e == _data._nil) return NULL; @@ -649,10 +644,12 @@ public: inline bool empty() const { return _data.size_cache == 0; } inline int size() const { return _data.size_cache; } + int calculate_depth() const { // used for debug mostly if (!_data._root) return 0; + int max_d = 0; _calculate_depth(_data._root->left, max_d, 0); return max_d; @@ -662,10 +659,10 @@ public: if (!_data._root) return; + _cleanup_tree(_data._root->left); _data._root->left = _data._nil; _data.size_cache = 0; - _data._nil->parent = _data._nil; _data._free_root(); }