Commit bb5614b

@@ -1,26 +1,175 @@
 #include "avl_tree.h"
+#include <stdio.h>
 #include <stdlib.h>
 
-AVLNode *avl_node_init(int value) {
-  AVLNode *node = malloc(sizeof(AVLNode));
+int max(int a, int b) {
+  return (a > b) ? a : b;
+}
+
+int height_of(AVLTree *node) {
+  return node == NULL ? 0 : node->height;
+}
+
+AVLTree *smallest(AVLTree *node) {
+  AVLTree *current = node;
+
+  while (current->left != NULL)
+    current = current->left;
+
+  return current;
+}
+
+AVLTree *rotate_right(AVLTree *y) {
+  AVLTree *x = y->left;
+  AVLTree *T2 = x->right;
+
+  x->right = y;
+  y->left = T2;
+
+  y->height = max(height_of(y->left), height_of(y->right)) + 1;
+  x->height = max(height_of(x->left), height_of(x->right)) + 1;
+
+  return x;
+}
+
+AVLTree *rotate_left(AVLTree *x) {
+  AVLTree *y = x->right;
+  AVLTree *T2 = y->left;
+
+  y->left = x;
+  x->right = T2;
+
+  x->height = max(height_of(x->left), height_of(x->right)) + 1;
+  y->height = max(height_of(y->left), height_of(y->right)) + 1;
+
+  return y;
+}
+
+int balance_of(AVLTree *node) {
+  return (node == NULL) ? 0 : height_of(node->left) - height_of(node->right);
+}
+
+AVLTree *avl_tree_initialize(int value) {
+  AVLTree *node = malloc(sizeof(AVLTree));
+  node->value = value;
   node->left = NULL;
   node->right = NULL;
-  node->value = value;
+  node->height = 1;
   return node;
 }
 
-AVLTree *avl_tree_init() {
-  AVLTree *tree = malloc(sizeof(AVLTree));
-  return tree;
+int avl_tree_size(AVLTree *tree) {
+  int total = 0;
+  if (tree == NULL)
+    return total;
+  if (tree->left)
+    total += avl_tree_size(tree->left);
+  if (tree->right)
+    total += avl_tree_size(tree->right);
+  return total + 1;
 }
 
-int avl_tree_size(AVLTree *tree) {
-  if (tree->root)
-    return 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 0;
+    return tree;
+
+  tree->height = 1 + max(height_of(tree->left), height_of(tree->right));
+
+  int balance = balance_of(tree);
+  if (balance > 1 && value < tree->left->value)
+    return rotate_right(tree);
+
+  if (balance < -1 && value > tree->right->value)
+    return rotate_left(tree);
+
+  if (balance > 1 && value > tree->left->value) {
+    tree->left = rotate_left(tree->left);
+    return rotate_right(tree);
+  }
+
+  if (balance < -1 && value < tree->right->value) {
+    tree->right = rotate_right(tree->right);
+    return rotate_left(tree);
+  }
+
+  return tree;
+}
+
+AVLTree *node_delete(AVLTree *root, int value) {
+  if (root == NULL)
+    return root;
+
+  if (value < root->value)
+    root->left = node_delete(root->left, value);
+  else if (value > root->value)
+    root->right = node_delete(root->right, value);
+  else {
+    if ((root->left == NULL) || (root->right == NULL)) {
+      AVLTree *temp = root->left ? root->left : root->right;
+
+      if (temp == NULL) {
+        temp = root;
+        root = NULL;
+      } else {
+        *root = *temp;
+      }
+      free(temp);
+    } else {
+      AVLTree *temp = smallest(root->right);
+      root->value = temp->value;
+      root->right = node_delete(root->right, temp->value);
+    }
+  }
+
+  if (root == NULL)
+    return root;
+
+  root->height = 1 + max(height_of(root->left), height_of(root->right));
+
+  int balance = balance_of(root);
+  if (balance > 1 && balance_of(root->left) >= 0)
+    return rotate_right(root);
+
+  if (balance > 1 && balance_of(root->left) < 0) {
+    root->left = rotate_left(root->left);
+    return rotate_right(root);
+  }
+
+  if (balance < -1 && balance_of(root->right) <= 0)
+    return rotate_left(root);
+
+  if (balance < -1 && balance_of(root->right) > 0) {
+    root->right = rotate_right(root->right);
+    return rotate_left(root);
+  }
+
+  return root;
+}
+
+void print_tree(AVLTree *tree, int level) {
+  for (int i = 0; i < level; i++)
+    printf(" ");
+
+  if (tree) {
+    printf("(%d)\n", tree->value);
+
+    if (!tree->left && !tree->right)
+      return;
+    print_tree(tree->left, level + 1);
+    print_tree(tree->right, level + 1);
+  }
+  else {
+    printf("( )\n");
+  }
 }
 
-void avl_tree_insert(AVLTree *tree, int value) {
-  tree->root = avl_node_init(value);
+void avl_tree_inspect(AVLTree *tree) {
+  print_tree(tree, 0);
 }

@@ -1,13 +1,11 @@
 typedef struct node {
   struct node *left;
   struct node *right;
+  int height;
   int value;
-} AVLNode;
-
-typedef struct {
-  AVLNode *root;
 } AVLTree;
 
-AVLTree *avl_tree_init(void);
+AVLTree *avl_tree_initialize(int value);
 int avl_tree_size(AVLTree *tree);
-void avl_tree_insert(AVLTree *tree, int value);
+AVLTree *avl_tree_insert(AVLTree *tree, int value);
+void avl_tree_inspect(AVLTree *tree);

@@ -3,21 +3,63 @@
 #include <string.h>
 
 Ensure(initialize_returns_new_tree) {
-  AVLTree *tree = avl_tree_init();
+  AVLTree *tree = avl_tree_initialize(1);
   assert_that(tree, is_not_equal_to(NULL));
 }
 
 Ensure(size_returns_zero) {
-  AVLTree *tree = avl_tree_init();
+  AVLTree *tree = avl_tree_initialize(1);
 
-  assert_that(avl_tree_size(tree), is_equal_to(0));
+  assert_that(avl_tree_size(tree), is_equal_to(1));
 }
 
 Ensure(insert_changes_size) {
-  AVLTree *tree = avl_tree_init();
-  avl_tree_insert(tree, 33);
+  AVLTree *tree = avl_tree_initialize(5);
+  avl_tree_insert(tree, 4);
 
-  assert_that(avl_tree_size(tree), is_equal_to(1));
+  assert_that(avl_tree_size(tree), is_equal_to(2));
+}
+
+Ensure(insert_changes_height) {
+  AVLTree *tree = avl_tree_initialize(5);
+  tree = avl_tree_insert(tree, 4);
+
+  assert_that(tree->height, is_equal_to(2));
+  assert_that(tree->left->height, is_equal_to(1));
+}
+
+Ensure(insert_performs_a_left_rotation) {
+/*
+  (10)                  (20)
+    \                  /    \
+    (20)      ->    (10)    (30)
+      \
+      (30)
+*/
+  AVLTree *tree = avl_tree_initialize(10);
+  tree = avl_tree_insert(tree, 20);
+  tree = avl_tree_insert(tree, 30);
+
+  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_rotation) {
+/*
+     (30)            (20)
+     /              /   \
+   (20)     -->   (10)  (30)
+   /
+(10)
+*/
+  AVLTree *tree = avl_tree_initialize(30);
+  tree = avl_tree_insert(tree, 20);
+  tree = avl_tree_insert(tree, 10);
+
+  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() {
@@ -25,6 +67,9 @@ TestSuite *avl_tree_tests() {
   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;
 }