AVL Tree Insertion Example

avl_insertion.java

package jono.structures;


import java.lang.Math;
import jono.structures.Node;
import java.util.stream.Stream;

class AvlNode extends Node {
  int balance = 0;
  int height  = 0;
  AvlNode parent;
  AvlNode left, right;

  public AvlNode(int value) {
    super(value);
  }

  public AvlNode insert(AvlNode node) {
    if (this.value >= node.value) {
      if (this.left == null) {
        this.left = node;
        node.parent = this;
      } else {
        this.left.insert(node);
      }
    } else {
      if (this.right == null) {
        this.right = node;
        node.parent = this;
      } else {
        this.right.insert(node);
      }
    }

    AvlNode parent = node.parent.rebalance();

    return parent;
  }

  public AvlNode rebalance() {
    AvlNode root = this;

    this.setBalance();

    if (this.balance <= -2) {
      if (AvlNode.height(this.left.left) > AvlNode.height(this.left.right)) {
        root = this.rotateRight();
      } else {
        root = this.rotateLeftThenRight();
      }
    }

    if (this.balance >= 2) {
      if (AvlNode.height(this.right.right) > AvlNode.height(this.right.left)) {
        root = this.rotateLeft();
      } else {
        root = this.rotateRightThenLeft();
      }
    }

    // Rebalance parents recursivly
    if (this.parent != null) {
      root = this.parent.rebalance();
    }

    return root;
  }

  public AvlNode rotateLeft() {
    System.out.println(this.value + " will rotate Left");

    AvlNode a = this;
    AvlNode b = a.right;
    b.parent = a.parent;

    a.right = b.left;

    if (a.right != null)
      a.right.parent = a;

    b.left = a;
    a.parent = b;

    if (b.parent != null) {
      if (b.parent.right == a) {
        b.parent.right = b;
      } else {
        b.parent.left = b;
      }
    }

    a.setBalance();
    b.setBalance();

    return b;
  }

  public AvlNode rotateRight() {
    System.out.println(this.value + " will rotate Right");

    AvlNode a = this;
    AvlNode b = a.left;

    b.parent = a.parent;

    a.left = b.right;

    if (a.left != null)
        a.left.parent = a;

    b.right = a;
    a.parent = b;

    if (b.parent != null) {
        if (b.parent.right == a) {
            b.parent.right = b;
        } else {
            b.parent.left = b;
        }
    }

    a.setBalance();
    b.setBalance();

    return b;
  }

  public AvlNode rotateLeftThenRight() {
    System.out.println(this.value + " left right rotation");

    this.left = this.left.rotateLeft();
    return this.rotateRight();
  }

  public AvlNode rotateRightThenLeft() {
    System.out.println(this.value + "right left rotation");
    this.right = this.right.rotateRight();
    return this.rotateLeft();
  }

  public void setBalance() {
    this.updateHeight();

    int left = (this.left == null) ? 0 : this.left.height;
    int right = (this.right == null) ? 0 : this.right.height;

    this.balance = right - left;
  }

  private void updateHeight() {
    int left = (this.left == null) ? -1 : this.left.height;
    int right = (this.right == null) ? -1 : this.right.height;

    this.height = 1 + Math.max(left, right);
  }

  public static int height(AvlNode n) {
    if (n == null) {
      return -1;
    }

    return n.height;
  }

  public static void main(String... args) {
    AvlNode root;

    int[] numbers = Stream.of(args[0].split(",")).mapToInt(Integer::parseInt).toArray();

    for(int i = 0; i < numbers.length; i++) {
      int number = numbers[i];

      System.out.printf("==== Inserting %d ====%n", number);

      if (root == null) {
        root = new AvlNode(number);
      } else {
        root = root.insert(new AvlNode(number));
      }

      System.out.println("Root is " + root.value);
      // root.inOrder();
    }

    System.out.println("======");
    root.inOrder();
    System.out.println("======");

    System.out.println(root.height);
    System.out.println(root.balance);
  }

  public void inOrder() {
    if (this.left != null) { this.left.inOrder(); }
    System.out.println(this.value);
    if (this.right != null) { this.right.inOrder(); }
  }

}