From 814d88737e836826f42a5bf0c6b02ed53579f1e2 Mon Sep 17 00:00:00 2001 From: Timerix22 Date: Thu, 2 May 2024 02:40:33 +0500 Subject: [PATCH] RBTree iterator --- src/RBTree.hpp | 241 +++++++++++++++++++++++++++++++++++++------------ 1 file changed, 185 insertions(+), 56 deletions(-) diff --git a/src/RBTree.hpp b/src/RBTree.hpp index b7f2579..76dde77 100644 --- a/src/RBTree.hpp +++ b/src/RBTree.hpp @@ -4,9 +4,9 @@ #include #include -// template -typedef char* TKey; -typedef char* TVal; +// typedef char* TKey; +// typedef char* TVal; +template class RBTree { enum class Color { Red, Black @@ -20,11 +20,9 @@ class RBTree { Node* left = nullptr; Node* right = nullptr; - Node(TKey _key, TVal _val, Color _color, Node* _parent){ - key = _key; - value = _val; - color = _color; - parent = _parent; + Node(TKey _key, TVal _val, Color _color, Node* _parent) + : key(_key), value(_val), color(_color), parent(_parent) + { } ~Node(){ @@ -53,10 +51,93 @@ class RBTree { return nullptr; return parent->getSibling(); } + + // n should be not null + Node* getMinChild(){ + Node* n = this; + while(n->left != nullptr) + n = n->left; + return n; + } + + // n should be not null + Node* getMaxChild(){ + Node* n = this; + while(n->right != nullptr) + n = n->right; + return n; + } }; + template + struct TreeIterator : std::iterator { + Node* n; + + TreeIterator(TreeIterator const& src){ + n = src.n; + } + + TreeIterator(Node* ptr){ + n = ptr; + } + + bool operator!=(TreeIterator const& other) const { + return n == other.n; + } + + bool operator==(TreeIterator const& other) const { + return n == other.n; + } + + reference operator*() const { + if(n == nullptr) + throw "RBTree::TreeIterator::operator*() error: n == nullptr"; + return *((pointer)(void*)n); + } + + TreeIterator& operator++() { + if(n == nullptr) + return *this; + Node* prev_n = nullptr; + while(n->parent != nullptr && (n->right == nullptr || n->right == prev_n)){ + prev_n = n; + n = n->parent; + } + n = n->right; + return *this; + } + + TreeIterator& operator--() { + if(n == nullptr) + return *this; + Node* prev_n = nullptr; + while(n->parent != nullptr && (n->left == nullptr || n->left == prev_n)){ + prev_n = n; + n = n->parent; + } + n = n->left; + return *this; + } + }; + + Node* root = nullptr; + ///@returns null if root is null + Node* findParentForKey(TKey key) const { + Node* n = root; + Node* parent = nullptr; + while(n != nullptr){ + parent = n; + if(key < n->key) + n = n->left; + else if(key > n->key) + n = n->right; + else return nullptr; // key == n->key + } + return parent; + } + void rotateLeft(Node* x){ // 1. get right child of x Node* y = x->right; @@ -96,19 +177,13 @@ class RBTree { x->parent = y; } - ///@returns null if root is null - Node* findParentForKey(TKey key){ - Node* n = root; - Node* parent = nullptr; - while(n != nullptr){ - parent = n; - if(key < n->key) - n = n->left; - else if(key > n->key) - n = n->right; - else return nullptr; // key == n->key - } - return parent; + void transplantNode(Node* old, Node* neww){ + neww->parent = old->parent; + if(old->parent == nullptr) + root = neww; + else if(old->parent->left == old) + old->parent->left = neww; + else old->parent->right = neww; } void fixupInsertion(Node* n){ @@ -221,41 +296,24 @@ class RBTree { } } - void transplantNode(Node* old, Node* neww){ - neww->parent = old->parent; - if(old->parent == nullptr) - root = neww; - else if(old->parent->left == old) - old->parent->left = neww; - else old->parent->right = neww; - } - - // n should be not null - Node* getMinimalChild(Node* n){ - while(n->left != nullptr) - n = n->left; - return n; - } - - // n should be not null - Node* getMaximalChild(Node* n){ - while(n->right != nullptr) - n = n->right; - return n; - } public: + using iterator = TreeIterator>; + using const_iterator = TreeIterator>; + RBTree() {} ~RBTree(){ delete root; } - - bool tryAdd(TKey key, TVal& value, TVal*& result){ + /// @param resultPtr nullable + bool tryAdd(TKey key, TVal& value, TVal** resultPtr){ if(root == nullptr){ root = new Node(key, value, Color::Black, nullptr); + if(resultPtr) + *resultPtr = &root->value; return true; } @@ -267,20 +325,35 @@ public: else nodePtrPtr = &parent->right; // if a child node already exists at this place, returns false - if(*nodePtrPtr != nullptr) + if(*nodePtrPtr != nullptr){ + if(resultPtr) + *resultPtr = nullptr; return false; + } // places newNode to left or right of the parent Node* newNode = new Node(key, value, Color::Red, parent); + if(resultPtr) + *resultPtr = &newNode->value; *nodePtrPtr = newNode; // auto-balancing fixupInsertion(newNode); return true; } + + /// @param resultPtr nullable + bool tryAdd(TKey key, TVal&& value, TVal** resultPtr){ + return tryAdd(key, value, resultPtr); + } - bool trySet(TKey key, TVal& value){ - if(root == nullptr) + + /// @param resultPtr nullable + bool trySet(TKey key, TVal& value, TVal** resultPtr){ + if(root == nullptr){ + if(resultPtr) + *resultPtr = nullptr; return false; + } Node* parent = findParentForKey(key); // ptr to parent->right or parent->left @@ -290,17 +363,31 @@ public: else nodePtrPtr = &parent->right; // if a child node with the given key doesn't exist, returns false - if(*nodePtrPtr == nullptr) + if(*nodePtrPtr == nullptr){ + if(resultPtr) + *resultPtr = nullptr; return false; + } // replaces the value of left or right child of the parent (*nodePtrPtr)->value = value; + if(resultPtr) + *resultPtr = &(*nodePtrPtr)->value; return true; } + + /// @param resultPtr nullable + bool trySet(TKey key, TVal&& value, TVal** resultPtr){ + return trySet(key, value, resultPtr); + } - void addOrSet(TKey key, TVal& value){ + + /// @param resultPtr nullable + void addOrSet(TKey key, TVal& value, TVal** resultPtr){ if(root == nullptr){ root = new Node(key, value, Color::Black, nullptr); + if(resultPtr != nullptr) + *resultPtr = &root->value; return; } @@ -314,38 +401,55 @@ public: // if a child node already exists at this place, sets it's value if(*nodePtrPtr != nullptr){ (*nodePtrPtr)->value = value; + if(resultPtr) + *resultPtr = &(*nodePtrPtr)->value; return; } // places newNode to left or right of the parent Node* newNode = new Node(key, value, Color::Red, parent); + if(resultPtr != nullptr) + *resultPtr = &newNode->value; *nodePtrPtr = newNode; // auto-balancing fixupInsertion(newNode); } - bool tryGet(TKey key, TVal** result){ + /// @param resultPtr nullable + void addOrSet(TKey key, TVal&& value, TVal** resultPtr){ + addOrSet(key, value, resultPtr); + } + + + bool tryGet(TKey key, TVal** resultPtr) const { + if(!resultPtr) + return false; Node* parent = findParentForKey(key); Node* n; if(key < parent->key) n = parent->left; else n = parent->right; // if there is no node with the given key - if(n == nullptr) + if(n == nullptr){ + *resultPtr = nullptr; return false; + } - *result = &n->value; + *resultPtr = &n->value; return true; } + bool tryDelete(TKey key){ Node* parent = findParentForKey(key); Node* n; if(key < parent->key) n = parent->left; else n = parent->right; - if(n == nullptr) // key not found + // key not found + if(n == nullptr){ return false; + } if(n->left == nullptr){ transplantNode(n, n->right); @@ -358,7 +462,7 @@ public: fixupDeletion(n->left); } else { - Node* minNode = getMinimalChild(n->right); + Node* minNode = n->right->getMinChild(); if(minNode != n->right){ transplantNode(minNode, minNode->right); minNode->right = n->right; @@ -373,5 +477,30 @@ public: } delete n; - } + } + + + iterator begin(){ + if(root == nullptr) + return iterator(nullptr); + return iterator(root->getMinChild()); + } + + iterator end(){ + if(root == nullptr) + return iterator(nullptr); + return iterator(root->getMaxChild()); + } + + const_iterator begin() const { + if(root == nullptr) + return const_iterator(nullptr); + return const_iterator(root->getMinChild()); + } + + const_iterator end() const { + if(root == nullptr) + return const_iterator(nullptr); + return const_iterator(root->getMaxChild()); + } };