AVL Binary Search Tree Implementation in Java

AVL (Adelson-Velsky and Landis) is an example of self-balanced Binary Search Tree. It has the same basic operations as that of a regular BST but modifications are followed by zero or more tree rotation operations. Every node also stores balance factor which is corrected after insertion or deletion. If the recomputed balance factor remains in the range from -1 to +1 then only corrections of the balance factor is required and no rotations are necessary. But if the recomputed balance factor becomes less than -1 or greater than +1 the subtree rooted at this node is unbalanced, and a rotation is needed. The AVLNode class (extends Node from BinarySearchTree class) can be declared like:

public static class AVLNode<T extends Comparable<T>> extends Node<T> {
    protected int height = 1;
    protected AVLNode(Node<T> parent, T value) {
        super(parent, value);
    }

    protected boolean isLeaf() {
        return ((getLeft() == null) && (getRight() == null));
    }

    protected void updateHeight() {
        int lesserHeight = 0;
        int greaterHeight = 0;
        if (getLeft() != null) {
            AVLNode<T> lesserAVLNode = (AVLNode<T>) getLeft();
            lesserHeight = lesserAVLNode.height;
        }
        if (getRight() != null) {
            AVLNode<T> greaterAVLNode = (AVLNode<T>) getRight();
            greaterHeight = greaterAVLNode.height;
        }
        if (lesserHeight > greaterHeight) {
            height = lesserHeight + 1;
        } else {
            height = greaterHeight + 1;
        }
    }

    protected int getBalanceFactor() {
        int lesserHeight = 0;
        int greaterHeight = 0;
        if (getLeft() != null) {
            AVLNode<T> lesserAVLNode = (AVLNode<T>) getLeft();
            lesserHeight = lesserAVLNode.height;
        }
        if (getRight() != null) {
            AVLNode<T> greaterAVLNode = (AVLNode<T>) getRight();
            greaterHeight = greaterAVLNode.height;
        }
        return greaterHeight - lesserHeight;
    }

    @Override
    public String toString() {
        return "value=" + getData() + " height=" + height + " parent=" + ((getParent() != null) ? getParent().getData() : "NULL")
                    + " left=" + ((getLeft() != null) ? getLeft().getData() : "NULL") + " right="
                    + ((getRight() != null) ? getRight().getData() : "NULL");
    }
}

The AVLTree class can be declared like:

public class AVLTree<T extends Comparable<T>> extends RotatableBinarySearchTree<T> {
   public AVLTree() { }
}

We need to re-balance after insertion and deletion. Consider the insertion method first

@Override
protected Node<T> insertValue(T id) {
        Node<T> returnedNode = super.insertValue(id);
 AVLNode<T> insertedNode = (AVLNode<T>) returnedNode;
 while (insertedNode != null) {
  insertedNode.updateHeight();
  balanceAfterInsert(insertedNode);
  insertedNode = (AVLNode<T>) insertedNode.getParent();
 }
 return returnedNode;
}

private enum RotationSequence {
        LEFT_LEFT, LEFT_RIGHT, RIGHT_LEFT, RIGHT_RIGHT
};

private void balanceAfterInsert(AVLNode<T> node) {
 int balanceFactor = node.getBalanceFactor();
 
 if (balanceFactor > 1 || balanceFactor < -1) {
  AVLNode<T> parent = null;
  AVLNode<T> child = null;
  RotationSequence rotation = null;
  
  if (balanceFactor < 0) {
   parent = (AVLNode<T>) node.getLeft();
   balanceFactor = parent.getBalanceFactor();
   if (balanceFactor < 0) {
    child = (AVLNode<T>) parent.getLeft();
    rotation = RotationSequence.LEFT_LEFT;
   } else {
    child = (AVLNode<T>) parent.getRight();
    rotation = RotationSequence.LEFT_RIGHT;
   }
  } else {
   parent = (AVLNode<T>) node.getRight();
   balanceFactor = parent.getBalanceFactor();
   if (balanceFactor < 0) {
    child = (AVLNode<T>) parent.getLeft();
    rotation = RotationSequence.RIGHT_LEFT;
   } else {
    child = (AVLNode<T>) parent.getRight();
    rotation = RotationSequence.RIGHT_RIGHT;
   }
  }

  if (rotation == RotationSequence.LEFT_RIGHT) {
   // Left-Right (Left rotation, right rotation)
   rotateLeft(parent);
   rotateRight(node);
  } else if (rotation == RotationSequence.RIGHT_LEFT) {
   // Right-Left (Right rotation, left rotation)
   rotateRight(parent);
   rotateLeft(node);
  } else if (rotation == RotationSequence.LEFT_LEFT) {
   // Left-Left (Right rotation)
   rotateRight(node);
  } else {
   // Right-Right (Left rotation)
   rotateLeft(node);
  }

  node.updateHeight(); // New child node
  child.updateHeight(); // New child node
  parent.updateHeight(); // New Parent node
 }
}

The deletion will be similar and the process is: find the node in tree, find its replacement and replace the node and balance it accordingly.

@Override
protected Node<T> deleteValue(T value) {
 // Find node to remove
 Node<T> nodeToRemoved = this.searchAndGetNode(value);
 if (nodeToRemoved != null) {
  // Find the replacement node
  Node<T> replacementNode = this.getReplacementNode(nodeToRemoved);

  // Find the parent of the replacement node to re-factor the
  // height/balance of the tree
  AVLNode<T> nodeToRefactor = null;
  if (replacementNode != null)
   nodeToRefactor = (AVLNode<T>) replacementNode.getParent();
  if (nodeToRefactor == null)
   nodeToRefactor = (AVLNode<T>) nodeToRemoved.getParent();
  if (nodeToRefactor != null && nodeToRefactor.equals(nodeToRemoved))
   nodeToRefactor = (AVLNode<T>) replacementNode;

  // Replace the node
  deleteAndReplaceNode(nodeToRemoved, replacementNode);

  // Re-balance the tree all the way up the tree
  if (nodeToRefactor != null) {
   while (nodeToRefactor != null) {
    nodeToRefactor.updateHeight();
    balanceAfterDelete(nodeToRefactor);
    nodeToRefactor = (AVLNode<T>) nodeToRefactor.getParent();
   }
  }
 }
 return nodeToRemoved;
}

private void balanceAfterDelete(AVLNode<T> node) {
 int balanceFactor = node.getBalanceFactor();
 if (balanceFactor == -2 || balanceFactor == 2) {
  if (balanceFactor == -2) {
   AVLNode<T> ll = (AVLNode<T>) node.getLeft().getLeft();
   int lesser = (ll != null) ? ll.height : 0;
   AVLNode<T> lr = (AVLNode<T>) node.getLeft().getRight();
   int greater = (lr != null) ? lr.height : 0;
   if (lesser >= greater) {
    rotateRight(node);
    node.updateHeight();
    if (node.getParent() != null)
     ((AVLNode<T>) node.getParent()).updateHeight();
   } else {
    rotateLeft(node.getLeft());
    rotateRight(node);

    AVLNode<T> p = (AVLNode<T>) node.getParent();
    if (p.getLeft() != null)
     ((AVLNode<T>) p.getLeft()).updateHeight();
    if (p.getRight() != null)
     ((AVLNode<T>) p.getRight()).updateHeight();
    p.updateHeight();
   }
  } else if (balanceFactor == 2) {
   AVLNode<T> rr = (AVLNode<T>) node.getRight().getRight();
   int greater = (rr != null) ? rr.height : 0;
   AVLNode<T> rl = (AVLNode<T>) node.getRight().getLeft();
   int lesser = (rl != null) ? rl.height : 0;
   if (greater >= lesser) {
    rotateLeft(node);
    node.updateHeight();
    if (node.getParent() != null)
     ((AVLNode<T>) node.getParent()).updateHeight();
   } else {
    rotateRight(node.getRight());
    rotateLeft(node);

    AVLNode<T> p = (AVLNode<T>) node.getParent();
    if (p.getLeft() != null)
     ((AVLNode<T>) p.getLeft()).updateHeight();
    if (p.getRight() != null)
     ((AVLNode<T>) p.getRight()).updateHeight();
    p.updateHeight();
   }
  }
 }
}

Remaining methods are inherited from BinarySearchTree class. So the major changes are in insertion/deletion along with re-balancing efforts. But considering the benefits it provide it is worth it.