fix: repaint colour when unbalanced

4 files changed, 245 insertions, 10 deletions

diff --git a/src/03/avl_tree.c b/src/03/avl_tree.c
index e1e3e66..bafe319 100644
--- a/src/03/avl_tree.c
+++ b/src/03/avl_tree.c
@@ -163,7 +163,7 @@ AVLTree *avl_tree_delete(AVLTree *tree, int value) {
   return tree;
 }
 
-void print_tree(AVLTree *tree, int level) {
+static void print_tree(AVLTree *tree, int level) {
   for (int i = 0; i < level; i++)
     printf(" ");
 
diff --git a/src/03/rb_tree.c b/src/03/rb_tree.c
index 4185da9..3ce001f 100644
--- a/src/03/rb_tree.c
+++ b/src/03/rb_tree.c
@@ -2,25 +2,178 @@
 #include <stdlib.h>
 #include <stdio.h>
 
+/*
+ * Implementation derived from:
+ * * https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion
+ */
+
+void insert_repair_tree(RBTree *tree);
+
+RBTree *parent_of(RBTree *node) {
+  return node ? node->parent : NULL;
+}
+
+RBTree *grand_parent_of(RBTree *node) {
+  return parent_of(parent_of(node));
+}
+
+RBTree *sibling_of(RBTree *node) {
+  RBTree *parent = parent_of(node);
+
+  if (!parent)
+    return NULL;
+
+  return node == parent->left ? parent->right : parent->left;
+}
+
+RBTree *pibling_of(RBTree *node) {
+  return sibling_of(parent_of(node));
+}
+
+static void rb_rotate_left(RBTree *tree) {
+  RBTree *tmp = tree->right;
+  RBTree *parent = parent_of(tree);
+
+  tree->right = tmp->left;
+  tmp->left = tree;
+  tree->parent = tmp;
+
+  if (tree->right)
+    tree->right->parent = tree;
+
+  if (parent) {
+    if (tree == parent->left)
+      parent->left = tmp;
+    else if (tree == parent->right)
+      parent->right = tmp;
+  }
+  tmp->parent = parent;
+}
+
+static void rb_rotate_right(RBTree *tree) {
+  RBTree *tmp = tree->left;
+  RBTree *parent = parent_of(tree);
+
+  tree->left = tmp->right;
+  tmp->right = tree;
+  tree->parent = tmp;
+
+  if (tree->left)
+    tree->left->parent = tree;
+
+  if (parent) {
+    if (tree == parent->left)
+      parent->left = tmp;
+    else if (tree == parent->right)
+      parent->right = tmp;
+  }
+  tmp->parent = parent;
+}
+
+void insert_case_4_step_2(RBTree *tree) {
+  RBTree *parent = parent_of(tree);
+  RBTree *grand_parent = grand_parent_of(tree);
+
+  if (tree == parent->left)
+    rb_rotate_right(grand_parent);
+  else
+    rb_rotate_left(grand_parent);
+  parent->colour = black;
+  grand_parent->colour = red;
+}
+
+void insert_repair_tree(RBTree *tree) {
+  RBTree *parent = parent_of(tree);
+  RBTree *pibling = pibling_of(tree);
+
+  if (parent == NULL || parent->colour == black) {
+    return;
+  } else if (pibling && pibling->colour == red) {
+    parent->colour = black;
+    pibling->colour = black;
+    grand_parent_of(tree)->colour = red;
+    insert_repair_tree(grand_parent_of(tree));
+  } else {
+    RBTree *grand_parent = grand_parent_of(tree);
+
+    if (tree == parent->right && parent == grand_parent->left) {
+      rb_rotate_left(parent);
+    } else if (tree == parent->left && parent == grand_parent->right) {
+      rb_rotate_right(parent);
+      tree = tree->right;
+    }
+
+    insert_case_4_step_2(tree);
+  }
+}
+
+static void insert(RBTree *root, RBTree *node) {
+  if (root) {
+    if (node ->value < root->value) {
+      if (root->left) {
+        insert(root->left, node);
+        return;
+      } else {
+        root->left = node;
+      }
+    } else {
+      if (root->right) {
+        insert(root->right, node);
+        return;
+      } else {
+        root->right = node;
+      }
+    }
+  }
+
+  node->parent = root;
+  node->left = NULL;
+  node->right = NULL;
+  node->colour = red;
+}
+
+static void print_tree(RBTree *tree, int level) {
+  for (int i = 0; i < level; i++)
+    printf(" ");
+
+  if (tree) {
+    printf("(%d:%c)\n", tree->value, tree->colour == red ? 'R' : 'B');
+
+    if (!tree->left && !tree->right)
+      return;
+    print_tree(tree->left, level + 1);
+    print_tree(tree->right, level + 1);
+  }
+  else {
+    printf("( )\n");
+  }
+}
+
 RBTree *rb_tree_initialize(int value) {
   RBTree *tree = malloc(sizeof(RBTree));
   tree->colour = black;
   tree->left = NULL;
+  tree->parent = NULL;
   tree->right = NULL;
   tree->value = value;
   return tree;
 }
 
 RBTree *rb_tree_insert(RBTree *tree, int value) {
+  RBTree *node = rb_tree_initialize(value);
+
   if (tree == NULL)
-    return rb_tree_initialize(value);
+    return node;
 
-  if (value < tree->value) {
-    tree->left = rb_tree_insert(tree->left, value);
-  } else if (value > tree->value) {
-    tree->right = rb_tree_insert(tree->right, value);
-  } else {
-    printf("KABOOM");
-  }
-  return tree;
+  insert(tree, node);
+  insert_repair_tree(node);
+
+  RBTree *root = node;
+  while (parent_of(root))
+    root = parent_of(root);
+  return root;
+}
+
+void rb_tree_inspect(RBTree *tree) {
+  print_tree(tree, 0);
 }
diff --git a/src/03/rb_tree.h b/src/03/rb_tree.h
index f818232..484a332 100644
--- a/src/03/rb_tree.h
+++ b/src/03/rb_tree.h
@@ -5,6 +5,7 @@ enum Colour {
 
 typedef struct rb_node {
   struct rb_node *left;
+  struct rb_node *parent;
   struct rb_node *right;
   enum Colour colour;
   int value;
@@ -12,3 +13,4 @@ typedef struct rb_node {
 
 RBTree *rb_tree_initialize(int value);
 RBTree *rb_tree_insert(RBTree *tree, int value);
+void rb_tree_inspect(RBTree *tree);
diff --git a/src/03/rb_tree_test.c b/src/03/rb_tree_test.c
index 04af925..5202454 100644
--- a/src/03/rb_tree_test.c
+++ b/src/03/rb_tree_test.c
@@ -68,6 +68,83 @@ Ensure(insert_adds_a_new_item_to_left_subtree) {
   assert_that(tree->left->value, is_equal_to(10));
 }
 
+Ensure(rb_tree_insert_performs_a_right_rotation) {
+/*
+      (30)          (20:b)
+      /            /    \
+   (20)     ->  (10:r)    (30:r)
+   /
+(10)
+
+*/
+  RBTree *tree = rb_tree_initialize(30);
+
+  tree = rb_tree_insert(tree, 20);
+  tree = rb_tree_insert(tree, 10);
+
+  assert_that(tree, is_not_equal_to(NULL));
+  assert_that(tree->value, is_equal_to(20));
+  assert_that(tree->colour, is_equal_to(black));
+
+  assert_that(tree->left->value, is_equal_to(10));
+  assert_that(tree->left->colour, is_equal_to(red));
+
+  assert_that(tree->right->value, is_equal_to(30));
+  assert_that(tree->right->colour, is_equal_to(red));
+}
+
+Ensure(rb_tree_insert_performs_a_left_rotation) {
+/*
+(10)                 (20:b)
+   \                 /    \
+   (20)     ->  (10:r)    (30:r)
+     \
+     (30)
+*/
+
+  RBTree *tree = rb_tree_initialize(10);
+  tree = rb_tree_insert(tree, 20);
+  tree = rb_tree_insert(tree, 30);
+
+  assert_that(tree, is_not_equal_to(NULL));
+  assert_that(tree->value, is_equal_to(20));
+  assert_that(tree->colour, is_equal_to(black));
+  assert_that(tree->left->value, is_equal_to(10));
+  assert_that(tree->left->colour, is_equal_to(red));
+  assert_that(tree->right->value, is_equal_to(30));
+  assert_that(tree->right->colour, is_equal_to(red));
+}
+
+Ensure(rb_tree_insert_repaints_the_new_node) {
+/*
+     (20:b)                   (20:r)
+     /    \                  /     \
+  (10:r)  (30:r)  -->   (10:b)    (30:b)
+  /                     /
+(5:r)                 (5:r)
+*/
+
+  RBTree *tree = rb_tree_initialize(20);
+  tree = rb_tree_insert(tree, 10);
+  tree = rb_tree_insert(tree, 30);
+  tree = rb_tree_insert(tree, 5);
+
+  rb_tree_inspect(tree);
+
+  assert_that(tree, is_not_equal_to(NULL));
+  assert_that(tree->value, is_equal_to(20));
+  assert_that(tree->colour, is_equal_to(red));
+
+  assert_that(tree->left->value, is_equal_to(10));
+  assert_that(tree->left->colour, is_equal_to(black));
+
+  assert_that(tree->left->left->value, is_equal_to(5));
+  assert_that(tree->left->left->colour, is_equal_to(red));
+
+  assert_that(tree->right->value, is_equal_to(30));
+  assert_that(tree->right->colour, is_equal_to(black));
+}
+
 TestSuite *rb_tree_tests() {
   TestSuite *x = create_test_suite();
   add_test(x, one_equals_one);
@@ -75,5 +152,8 @@ TestSuite *rb_tree_tests() {
   add_test(x, insert_returns_a_new_tree_when_null);
   add_test(x, insert_adds_a_new_item_to_right_subtree);
   add_test(x, insert_adds_a_new_item_to_left_subtree);
+  add_test(x, rb_tree_insert_performs_a_right_rotation);
+  add_test(x, rb_tree_insert_performs_a_left_rotation);
+  add_test(x, rb_tree_insert_repaints_the_new_node);
   return x;
 }