Red black tree

rbtree.cc

#include <iostream>
#include <queue>

using namespace std;

enum Color { RED, BLACK };

struct Node {
  char val;
  Color color;
  Node *left;
  Node *right;
  Node *parent;
  Node(char v)
      : val(v), color(RED), left(nullptr), right(nullptr), parent(nullptr){};
  Node(char v, Color c)
      : val(v), color(c), left(nullptr), right(nullptr), parent(nullptr){};
};

// 4 scenarios
// 1. root => recolor itself
// 2. uncle = red => rotate parent
// 3. uncle = black, triangle => rotate parent
// 4. uncle = black, line => rotate grandparent, recolor parent / grandparent
void insert(Node *root, char val) {
  if (root == nullptr) {
    root = new Node(val, BLACK);
    return;
  }
  Node *newNode = new Node(val);
  Node *curr = root;
  Node *parent = nullptr;
  while (curr != nullptr) {
    parent = curr;
    if (curr->val >= val) {
      curr = curr->right;
    } else {
      curr = curr->left;
    }
  }
  newNode->parent = parent;
  if (parent->val >= val) {
    parent->right = newNode;
  } else {
    parent->left = newNode;
  }
  adjustRbt(root, newNode);
}

void adjustRbt(Node* root, Node *current) {
  while (current->parent->color == RED) {
    // 上半部：parent是grandparent的left child
    if (current->parent == current->parent->parent->left) {
      Node *uncle = current->parent->parent->right;
      // case1: 若uncle是紅色
      if (uncle->color == RED) {
        current->parent->color = BLACK;
        uncle->color = BLACK;
        current->parent->parent->color = RED; // grandparent改成紅色
        current = current->parent->parent;
      } else {
        // case2 & 3: uncle是黑色
        if (current == current->parent->right) { // case2
          current = current->parent;
          leftRotation(current);
        }
        // case3
        current->parent->color = BLACK;         //把parent塗黑
        current->parent->parent->color = RED; // grandparent塗紅
        rightRotation(current->parent->parent);
      }
    } else {
      // 下半部：parent是grandparent的right child, 與上半部對稱
      Node *uncle = current->parent->parent->left;
      // case1: 若uncle是紅色
      if (uncle->color == RED) {
        current->parent->color = BLACK;
        uncle->color = BLACK;
        current->parent->parent->color = RED; // grandparent改成紅色
        current = current->parent->parent;
      } else {
         // case2 & 3: uncle是黑色
        if (current == current->parent->left) { // case2
          current = current->parent;
          rightRotation(current);
        }
        // case3
        current->parent->color = BLACK;
        current->parent->parent->color = RED;
        leftRotation(current->parent->parent);
      }
    }
  }
  root->color = BLACK; // 確保root是黑色
}

//     x                 y
//       \             /
//         y   =>    x
//       /             \
//     z                 z
void leftRotation(Node *node) {
  Node *parent = node->parent;
  Node *grandPa = parent->parent;
  Node *grandGrandPa = grandPa->parent;
  if (grandGrandPa->val >= parent->val) {
    grandGrandPa->left = parent;
  } else {
    grandGrandPa->right = parent;
  }
  parent->parent = grandGrandPa;
  parent->left = grandPa;
  grandPa->parent = parent;
  grandPa->right = node;
  node->parent = grandPa;
}

//      y             x
//    /                 \
//  x         =>          y
//    \                 /
//      j             j
void rightRotation(Node *node) {
  Node *parent = node->parent;
  Node *grandPa = parent->parent;
  Node *grandGrandPa = parent->parent;
  if (grandGrandPa->val > parent->val) {
    grandGrandPa->left = parent;
  } else {
    grandGrandPa->right = parent;
  }
  parent->parent = grandGrandPa;
  parent->right = grandPa;
  grandPa->parent = parent;
  grandPa->left = node;
  node->parent = grandPa;
}

void remove(char val) {}

char getNodeVal(Node *node) {
  if (node != nullptr) {
    return node->val;
  }
  return ' ';
}

void inOrder(Node *node) {
  queue<Node *> q;
  q.push(node);

  while (!q.empty()) {
    const int size = q.size();
    for (int i = 0; i < size; i++) {
      Node *ptr = q.front();
      q.pop();
      cout << "curr: " << getNodeVal(ptr)
           << ", parent: " << getNodeVal(ptr->parent)
           << ", left: " << getNodeVal(ptr->left)
           << ", right: " << getNodeVal(ptr->right) << " ";
      if (ptr->left != nullptr) {
        q.push(ptr->left);
      }
      if (ptr->right != nullptr) {
        q.push(ptr->right);
      }
    }
    cout << endl;
  }
}
//       a
//      /
//     x
//   /   \
// i       y
//       /   \
//      j      k
int main() {
  Node *head = new Node('a');
  Node *x = new Node('x');
  Node *y = new Node('y');
  Node *i = new Node('i');
  Node *j = new Node('j');
  Node *k = new Node('k');
  head->left = x;
  x->left = i;
  x->right = y;
  y->left = j;
  y->right = k;
  x->parent = head;
  i->parent = x;
  y->parent = x;
  j->parent = y;
  k->parent = y;
  leftRotation(j);
  inOrder(head);
}