5 files changed, 68 insertions, 9 deletions

diff --git a/src/03/avl_tree.c b/src/03/avl_tree.c
index 9b5e23f..4637610 100644
--- a/src/03/avl_tree.c
+++ b/src/03/avl_tree.c
@@ -188,28 +188,54 @@ static bool is_odd(int value) {
   return !is_even(value);
 }
 
-RBTree *_avl_tree_to_rb_tree(AVLTree *tree, AVLTree *parent) {
+static RBTree *to_rb_tree(AVLTree *tree, AVLTree *parent) {
   if (!tree)
     return NULL;
 
-  enum Colour colour = (parent && is_even(parent->height) && is_odd(tree->height)) ? red : black;
-  RBTree *rb_tree = rb_tree_initialize_with(tree->value, colour);
+  RBTree *rb_tree = rb_tree_initialize(tree->value);
 
-  rb_tree->left = _avl_tree_to_rb_tree(tree->left, tree);
+  rb_tree->left = to_rb_tree(tree->left, tree);
   if (rb_tree->left)
     rb_tree->left->parent = rb_tree;
 
-  rb_tree->right = _avl_tree_to_rb_tree(tree->right, tree);
+  rb_tree->right = to_rb_tree(tree->right, tree);
   if (rb_tree->right)
     rb_tree->right->parent = rb_tree;
   return rb_tree;
 }
 
+static void colour_children(RBTree *a, RBTree *b);
+
+static void colour(RBTree* tree, enum Colour colour) {
+  if (!tree)
+    return;
+
+  tree->colour = colour;
+  colour_children(tree->left, tree->right);
+}
+
+static void colour_children(RBTree *a, RBTree *b) {
+  int a_height = rb_tree_height(a);
+  int b_height = rb_tree_height(b);
+
+  if (a_height < b_height || is_odd(a_height))
+    colour(a, black);
+  else
+    colour(a, red);
+
+  if (b_height < a_height || is_odd(b_height))
+    colour(b, black);
+  else
+    colour(b, red);
+}
+
 RBTree *avl_tree_to_rb_tree(AVLTree *tree) {
   if (!tree)
     return NULL;
 
-  return _avl_tree_to_rb_tree(tree, NULL);
+  RBTree *rb_tree = to_rb_tree(tree, NULL);
+  colour(rb_tree, black);
+  return rb_tree;
 }
 
 void avl_tree_inspect(AVLTree *tree) {
diff --git a/src/03/avl_tree_test.c b/src/03/avl_tree_test.c
index 56e412b..95e52a3 100644
--- a/src/03/avl_tree_test.c
+++ b/src/03/avl_tree_test.c
@@ -326,8 +326,11 @@ Ensure(to_rb_tree_returns_a_new_red_black_tree) {
     expected = rb_tree_insert(expected, items[i]);
   }
 
-  RBTree *rb_tree = avl_tree_to_rb_tree(tree);
-  assert_that(rb_equals(expected, rb_tree), is_equal_to(true));
+  RBTree *actual = avl_tree_to_rb_tree(tree);
+
+  assert_that(rb_equals(expected, actual), is_equal_to(true));
+  assert_that(rb_tree_is_valid(actual), is_equal_to(true));
+  assert_that(rb_tree_is_valid(expected), is_equal_to(true));
 }
 
 Ensure(to_rb_tree_handles_trees_with_a_large_depth) {
diff --git a/src/03/rb_tree.c b/src/03/rb_tree.c
index 802791e..b6d0700 100644
--- a/src/03/rb_tree.c
+++ b/src/03/rb_tree.c
@@ -176,7 +176,7 @@ static void print_tree(RBTree *tree, int level) {
     printf(" ");
 
   if (tree) {
-    printf("(%d:%c P:%d)\n", tree->value, tree->colour == red ? 'R' : 'B', tree->parent ? tree->parent->value : -1);
+    printf("(%d%c H:%d)\n", tree->value, tree->colour == red ? 'R' : 'B', rb_tree_height(tree));
 
     if (!tree->left && !tree->right)
       return;
@@ -261,3 +261,10 @@ bool rb_tree_is_valid(RBTree *tree) {
 
   return rb_tree_is_valid(tree->left) && rb_tree_is_valid(tree->right);
 }
+
+int rb_tree_height(RBTree *tree) {
+  if (!tree)
+    return 1;
+
+  return 1 + max(rb_tree_height(tree->left), rb_tree_height(tree->right));
+}
diff --git a/src/03/rb_tree.h b/src/03/rb_tree.h
index c5d9ffa..ca423a4 100644
--- a/src/03/rb_tree.h
+++ b/src/03/rb_tree.h
@@ -20,3 +20,4 @@ bool rb_equals(RBTree *tree, RBTree *other_tree);
 bool rb_tree_is_valid(RBTree *tree);
 int rb_tree_size(RBTree *tree);
 void rb_tree_inspect(RBTree *tree);
+int rb_tree_height(RBTree *tree);
diff --git a/src/03/rb_tree_test.c b/src/03/rb_tree_test.c
index 3fca748..4bb2d7e 100644
--- a/src/03/rb_tree_test.c
+++ b/src/03/rb_tree_test.c
@@ -277,6 +277,24 @@ Ensure(is_valid_return_true) {
   assert_that(rb_tree_is_valid(tree), is_equal_to(true));
 }
 
+Ensure(height_returns_one) {
+  assert_that(rb_tree_height(NULL), is_equal_to(1));
+}
+
+Ensure(height_returns_three_when_left_subtree_is_present) {
+  RBTree *tree = rb_tree_initialize(10);
+  tree = rb_tree_insert(tree, 5);
+
+  assert_that(rb_tree_height(tree), is_equal_to(3));
+}
+
+Ensure(height_returns_three_when_right_subtree_is_present) {
+  RBTree *tree = rb_tree_initialize(10);
+  tree = rb_tree_insert(tree, 15);
+
+  assert_that(rb_tree_height(tree), is_equal_to(3));
+}
+
 TestSuite *rb_tree_tests() {
   TestSuite *x = create_test_suite();
 
@@ -305,5 +323,9 @@ TestSuite *rb_tree_tests() {
   add_test(x, is_valid_returns_false_when_red_node_has_red_child);
   add_test(x, is_valid_returns_false_when_each_path_to_leaves_does_not_contain_the_same_number_of_black_nodes);
   add_test(x, is_valid_return_true);
+
+  add_test(x, height_returns_one);
+  add_test(x, height_returns_three_when_left_subtree_is_present);
+  add_test(x, height_returns_three_when_right_subtree_is_present);
   return x;
 }