add tests for left-right rotation and right-left rotation

2 files changed, 83 insertions, 19 deletions

diff --git a/src/03/avl_tree.c b/src/03/avl_tree.c
index d6d94d1..cc7d52f 100644
--- a/src/03/avl_tree.c
+++ b/src/03/avl_tree.c
@@ -69,19 +69,15 @@ int avl_tree_size(AVLTree *tree) {
   return total + 1;
 }
 
-AVLTree *avl_tree_insert(AVLTree *tree, int value) {
-  if (tree == NULL)
-    return avl_tree_initialize(value);
-
-  if (value < tree->value)
-    tree->left = avl_tree_insert(tree->left, value);
-  else if (value > tree->value)
-    tree->right = avl_tree_insert(tree->right, value);
-  else
-    return tree;
+int compare(int a, int b)
+{
+  if (a < b) return -1;
+  if (a > b) return 1;
 
-  tree->height = 1 + max(height_of(tree->left), height_of(tree->right));
+  return 0;
+}
 
+AVLTree *rebalance(AVLTree *tree, int value) {
   int balance = balance_of(tree);
   if (balance > 1 && value < tree->left->value)
     return rotate_right(tree);
@@ -102,6 +98,28 @@ AVLTree *avl_tree_insert(AVLTree *tree, int value) {
   return tree;
 }
 
+AVLTree *avl_tree_insert(AVLTree *tree, int value) {
+  if (tree == NULL)
+    return avl_tree_initialize(value);
+
+  switch(compare(value, tree->value)) {
+    case -1:
+      tree->left = avl_tree_insert(tree->left, value);
+      break;
+    case 1:
+      tree->right = avl_tree_insert(tree->right, value);
+      break;
+    default:
+      return tree;
+  }
+
+  tree->height = 1 + max(
+    height_of(tree->left),
+    height_of(tree->right)
+  );
+  return rebalance(tree, value);
+}
+
 AVLTree *node_delete(AVLTree *root, int value) {
   if (root == NULL)
     return root;
diff --git a/src/03/avl_tree_test.c b/src/03/avl_tree_test.c
index 3f7a0e8..ccfcefb 100644
--- a/src/03/avl_tree_test.c
+++ b/src/03/avl_tree_test.c
@@ -28,6 +28,15 @@ Ensure(insert_changes_height) {
   assert_that(tree->left->height, is_equal_to(1));
 }
 
+Ensure(insert_creates_a_new_root) {
+  AVLTree *tree = avl_tree_insert(NULL, 10);
+
+  assert_that(tree, is_not_equal_to(NULL));
+  assert_that(tree->value, is_equal_to(10));
+  assert_that(tree->left, is_equal_to(NULL));
+  assert_that(tree->right, is_equal_to(NULL));
+}
+
 Ensure(insert_performs_a_left_rotation) {
 /*
   (10)                  (20)
@@ -62,15 +71,52 @@ Ensure(insert_performs_a_right_rotation) {
   assert_that(tree->right->value, is_equal_to(30));
 }
 
+Ensure(insert_performs_a_left_right_rotation) {
+/*
+   (30)         (20)
+   /           /    \
+(10)     -> (10)    (30)
+   \
+   (20)
+*/
+  AVLTree *tree = avl_tree_initialize(30);
+  tree = avl_tree_insert(tree, 10);
+  tree = avl_tree_insert(tree, 20);
+
+  assert_that(tree->value, is_equal_to(20));
+  assert_that(tree->left->value, is_equal_to(10));
+  assert_that(tree->right->value, is_equal_to(30));
+}
+
+Ensure(insert_performs_a_right_left_rotation) {
+/*
+(10)             (20)
+    \            /  \
+    (30)  --> (10)  (30)
+   /
+ (20)
+*/
+  AVLTree *tree = avl_tree_initialize(10);
+  tree = avl_tree_insert(tree, 30);
+  tree = avl_tree_insert(tree, 20);
+
+  assert_that(tree->value, is_equal_to(20));
+  assert_that(tree->left->value, is_equal_to(10));
+  assert_that(tree->right->value, is_equal_to(30));
+}
+
 TestSuite *avl_tree_tests() {
-  TestSuite *suite = create_test_suite();
-  add_test(suite, initialize_returns_new_tree);
-  add_test(suite, size_returns_zero);
-  add_test(suite, insert_changes_size);
-  add_test(suite, insert_changes_height);
-  add_test(suite, insert_performs_a_left_rotation);
-  add_test(suite, insert_performs_a_right_rotation);
-  return suite;
+  TestSuite *x = create_test_suite();
+  add_test(x, initialize_returns_new_tree);
+  add_test(x, size_returns_zero);
+  add_test(x, insert_changes_size);
+  add_test(x, insert_changes_height);
+  add_test(x, insert_creates_a_new_root);
+  add_test(x, insert_performs_a_left_rotation);
+  add_test(x, insert_performs_a_right_rotation);
+  add_test(x, insert_performs_a_left_right_rotation);
+  add_test(x, insert_performs_a_right_left_rotation);
+  return x;
 }
 
 int main(int argc, char **argv) {