BUPT-homework/semester3/pset2/2-3-AVL-Tree.c

#include <stdio.h>
#include <stdlib.h>

typedef struct treeNode {
  int val;
  struct treeNode *l;
  struct treeNode *r;
  int height; // used to calculate balance factor
} treeNode;

int max(int a, int b) { return (a > b) ? a : b; }

treeNode *newNode(int val) {
  treeNode *root = malloc(sizeof(treeNode));
  root->val = val;
  root->l = NULL;
  root->r = NULL;
  root->height = 1; // update this manually
  return root;
}

int height(treeNode *tr) {
  // This function can be optimized.
  if (!tr) {
    return 0;
  } else {
    return 1 + max(height(tr->l), height(tr->r));
  }
}

treeNode *leftRotate(treeNode *x) {
  treeNode *y = x->r;
  x->r = y->l;
  y->l = x;

  // Update height here.
  x->height = height(x);
  y->height = height(y);

  return y;
}

treeNode *rightRotate(treeNode *y) {
  treeNode *x = y->l;
  y->l = x->r;
  x->r = y;

  x->height = height(x);
  y->height = height(y);

  return x;
}

int getBalance(treeNode *tr) {
  if (!tr) {
    return 0;
  } else {
    return height(tr->l) - height(tr->r);
  }
}

treeNode *insertTree(treeNode *root, int val) {
  // Use recursion.
  if (!root) {
    return newNode(val);
  } else {
    if (val > root->val) {
      root->r = insertTree(root->r, val);
    } else {
      root->l = insertTree(root->l, val);
    }

    root->height = height(root);

    int bal = getBalance(root);
    // check the balance using balance and comparing val with childs
    if (bal > 1) {
      if (val < root->l->val) {
        // left left
        return rightRotate(root);
      } else {
        // left right
        root->l = leftRotate(root->l);
        return rightRotate(root);
      }
    } else if (bal < -1) {
      if (val > root->r->val) {
        // right right
        return leftRotate(root);
      } else {
        // right left
        root->r = rightRotate(root->r);
        return leftRotate(root);
      }
    } else {
      // the balance is correct, don't change anything
      return root;
    }
  }
}

void freeTree(treeNode *root) {
  if (root) {
    freeTree(root->l);
    freeTree(root->r);
    free(root);
  }
}

void preorderTraversal(treeNode *root) {
  if (root) {
    printf("%d,", root->val);
    preorderTraversal(root->l);
    preorderTraversal(root->r);
  }
}

int main(void) {
  int tmp;
  treeNode *root = NULL;

  while (scanf("%d,", &tmp) != EOF) {
    // update every time, since root can be changed.
    root = insertTree(root, tmp);
  }

  preorderTraversal(root);

  freeTree(root);
  return 0;
}