RBTree iterator

src/RBTree.hpp

@@ -4,9 +4,9 @@
 #include <map>
 #include <memory>
 
-// template<typename TKey, typename TVal>
-typedef char* TKey;
-typedef char* TVal;
+// typedef char* TKey;
+// typedef char* TVal;
+template<typename TKey, typename TVal>
 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<typename TIteratorValue>
+    struct TreeIterator : std::iterator<std::bidirectional_iterator_tag, TIteratorValue> {
+        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,30 +296,11 @@ 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<std::pair<const TKey, TVal>>;
+    using const_iterator = TreeIterator<std::pair<const TKey, const TVal>>;
+
     RBTree() {}
 
     ~RBTree(){
@@ -252,10 +308,12 @@ public:
     }
 
 
-
-    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;
     }
     
-    bool trySet(TKey key, TVal& value){
-        if(root == nullptr)
+    /// @param resultPtr nullable
+    bool tryAdd(TKey key, TVal&& value, TVal** resultPtr){
+        return tryAdd(key, value, resultPtr);
+    }
+
+
+    /// @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;
     }
  
-    void addOrSet(TKey key, TVal& value){
+    /// @param resultPtr nullable
+    bool trySet(TKey key, TVal&& value, TVal** resultPtr){
+        return trySet(key, value, resultPtr);
+    }
+
+
+    /// @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;
@@ -374,4 +478,29 @@ 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());
+    }
 };