RBTree insertion

src/RBTree.hpp

@@ -4,35 +4,248 @@
 #include <map>
 #include <memory>
 
-template<typename TKey, typename TVal>
+// template<typename TKey, typename TVal>
+typedef char* TKey;
+typedef char* TVal;
 class RBTree {
     enum class Color {
-        Red, Black;
+        Red, Black
     };
 
     struct Node {
         TKey key;
         TVal value;
         Color color;
+        Node* parent;
+        Node* left = nullptr;
+        Node* right = nullptr;
+
+        Node(TKey _key, TVal _val, Color _color, Node* _parent){
+            key = _key;
+            value = _val;
+            color = _color;
+            parent = _parent;
+        }
+
+        ~Node(){
+            if(left != nullptr)
+                delete left;
+            if(right != nullptr)
+                delete right;
+        }
+
+        inline Node* getSibling(){
+            if(parent == nullptr)
+                return nullptr;
+            else if(parent->left == this)
+                return parent->right;
+            else return parent->left;
+        }
+
+        inline Node* getGrandparent(){
+            if(parent == nullptr)
+                return nullptr;
+            else return parent->parent;
+        }
+
+        inline Node* getUncle(){
+            if(parent == nullptr)
+                return nullptr;
+            return parent->getSibling();
+        }
     };
 
+    Node* root = nullptr;
+
+    void rotateLeft(Node* x){
+        // 1. get right child of x
+        Node* y = x->right;
+        // 2. move y to the position of x
+        y->parent = x->parent;
+        if (x->parent != nullptr){ // x != root
+            if(x == x->parent->left)
+                x->parent->left = y;
+            else x->parent->right = y;
+            // TODO: maybe should set root to y?
+        }
+        // 3. move y.left to x.right if it exists
+        x->right = y->left;
+        if (x->right != nullptr)
+            x->right->parent = x;
+        // 4. move x to y.left
+        y->left = x;
+        x->parent = y;
+    }
+
+    void rotateRight(Node* x){
+        // 1. get left child of x
+        Node* y = x->left;
+        // 2. move y up
+        y->parent = x->parent;
+        if (x->parent != nullptr){ // x != root
+            if(x == x->parent->left)
+                x->parent->left = y;
+            else x->parent->right = y;
+        }
+        // 3. move y.right to x.left if it exists
+        x->left = y->right;
+        if (x->left != nullptr)
+            x->left->parent = x;
+        // 4. move x to y.right
+        y->right = x;
+        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 fixupInsertion(Node* n){
+        // case 1: n is root -- root must be black
+        if (n->parent == nullptr){
+		    n->color = Color::Black;
+            return;
+        }
+
+        // case 2: parent is black -- no requirements mismatch
+        if (n->parent->color == Color::Black)
+		    return;
+
+        // case 3: parent and uncle are red -- red nodes must have black parents
+        Node* u = n->getUncle();
+        Node* g = n->getGrandparent();
+        if(u != nullptr && u->color == Color::Red){
+            n->parent->color = Color::Black;
+            u->color = Color::Black;
+            g->color = Color::Red;
+            fixupInsertion(g);
+            return;
+        }
+
+        // case 4: parent is red and uncle is black -- red nodes must have black parents
+        if ((n == n->parent->right) && (n->parent == g->left)) {
+            rotateLeft(n->parent);
+            n = n->left;
+        }
+        else if ((n == n->parent->left) && (n->parent == g->right)) {
+            rotateRight(n->parent);
+            n = n->right;
+        }
+
+        // case 5
+        n->parent->color = Color::Black;
+        g->color = Color::Red;
+        if ((n == n->parent->left) && (n->parent == g->left))
+            rotateRight(g);
+        else rotateLeft(g);
+    }
+    
+    void fixupDeletion(Node* n){
+        
+    }
+
+
 public:
     RBTree() {}
 
+    ~RBTree(){
+        delete root;
+    }
+
     bool tryAdd(TKey key, TVal& value, TVal*& result){
+        if(root == nullptr){
+            root = new Node(key, value, Color::Black, nullptr);
+            return true;
+        }
         
+        Node* parent = findParentForKey(key);
+        // ptr to parent->right or parent->left
+        Node** leftOrRightPtr;
+        if(key < parent->key)
+            leftOrRightPtr = &parent->left;
+        else leftOrRightPtr = &parent->right;
+            
+        // if a child node already exists at this place, returns false
+        if(*leftOrRightPtr != nullptr)
+            return false;
+        
+        // places newNode to left or right of the parent
+        Node* newNode = new Node(key, value, Color::Red, parent);
+        *leftOrRightPtr = newNode;
+        // auto-balancing
+        fixupInsertion(newNode);
+        return true;
     }
 
     bool trySet(TKey key, TVal& value){
+        if(root == nullptr)
+            return false;
         
+        Node* parent = findParentForKey(key);
+        // ptr to parent->right or parent->left
+        Node** leftOrRightPtr;
+        if(key < parent->key)
+            leftOrRightPtr = &parent->left;
+        else leftOrRightPtr = &parent->right;
+
+        // if a child node with the given key doesn't exist, returns false
+        if(*leftOrRightPtr == nullptr)
+            return false;
+        
+        // replaces the value of left or right child of the parent
+        (*leftOrRightPtr)->value = value;
+        return true;
     }
 
     void addOrSet(TKey key, TVal& value){
+        if(root == nullptr){
+            root = new Node(key, value, Color::Black, nullptr);
+            return;
+        }
         
+        Node* parent = findParentForKey(key);
+        // ptr to parent->right or parent->left
+        Node** leftOrRightPtr;
+        if(key < parent->key)
+            leftOrRightPtr = &parent->left;
+        else leftOrRightPtr = &parent->right;
+            
+        // if a child node already exists at this place, sets it's value
+        if(*leftOrRightPtr != nullptr){
+            (*leftOrRightPtr)->value = value;
+            return;
+        }
+        
+        // places newNode to left or right of the parent
+        Node* newNode = new Node(key, value, Color::Red, parent);
+        *leftOrRightPtr = newNode;
+        // auto-balancing
+        fixupInsertion(newNode);
     }
 
-    bool tryGet(TKey key, TVal*& result){
-
+    bool tryGet(TKey key, TVal** result){
+        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)
+            return false;
+        
+        *result = &n->value;
+        return true;
     }
 
     bool tryDelete(TKey key){