diff --git a/src/RBTree.hpp b/src/RBTree.hpp index 14b668f..4569d71 100644 --- a/src/RBTree.hpp +++ b/src/RBTree.hpp @@ -4,35 +4,248 @@ #include #include -template +// template +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){