AVL Tree - Height Balancing Binary Search Tree (BST) Implementation [Java]

This program implements AVL tree (written in Java). AVL trees are binary search trees which automatically balance their heights after insertion of elements or removal of existing elements.

A post on BSTs :

http://coderbots.blogspot.in/2014/09/binary-search-tree-operations-find.html

Source:

import java.util.Scanner;

public class AVL {

    Node root;

    class Node {

        Node father, left, right;
        int height = 0;
        int key;

        public Node(int key, Node left, Node right, Node father) {
            this.key = key;
            this.left = left;
            this.right = right;
            this.father = father;
        }

        public String toString() {
            return "key = " + key + " height = " + height;
        }
    }

    public void insert(Node toInsert, Node node) {
        if (root == null) {
            root = toInsert;
            return;
        }

        if (toInsert.key < node.key) {
            if (node.left != null) {
                insert(toInsert, node.left);
                return;
            } else {
                toInsert.father = node;
                node.left = toInsert;
            }
        } else if (toInsert.key >= node.key) {
            if (node.right != null) {
                insert(toInsert, node.right);
                return;
            } else {
                toInsert.father = node;
                node.right = toInsert;
            }
        }
        updateHeights(toInsert);
        balance(toInsert);
    }

    public Node findNode(Node findNode, Node node) {
        if (root == null) {
            return null;
        }

        if (findNode.key < node.key) {
            if (node.left != null) {
                return findNode(findNode, node.left);
            }
        } else if (findNode.key > node.key) {
            if (node.right != null) {
                return findNode(findNode, node.right);
            }
        } else if (findNode.key == node.key) {
            return node;
        }
        return null;
    }

    public boolean deleteNode(Node deleteNode, Node node) {
        if (root == null) {
            return false;
        } else if (root.key == deleteNode.key) {
            if (root.left == null && root.right == null) {
                root = null;
            } else if (root.left == null) {
                root = root.right;
                root.father = null;
            } else if (root.right == null) {
                root = root.left;
                root.father = null;
            } else {
                Node min = root.right;
                //min element has no left or right child.
                if (min.left == null && min.right == null) {
                    root.key = min.key;
                    root.right = null;
                    updateHeights(root);
                    balance(root);
                    return true;
                }
                while (min.left != null) {
                    //Min element has left child.
                    min = min.left;
                }
                root.key = min.key;
                if (min.right != null) {
                    min.right.father = min.father;
                }
                //Min element has no left child but has right child.
                if (min.father != root) {
                    min.father.left = min.right; 
                } else {
                    min.father.right = min.right; 
                }
                updateHeights(min.father);
                balance(min.father);
            }
            return true;
        } else {
            Node locatedNode = findNode(deleteNode, root);
            if (locatedNode != null) {
                deleteHelper(locatedNode);
            } else {
                return false;
            }
        }
        return true;
    }

    private void deleteHelper(Node locatedNode) {
        boolean left = (locatedNode.father.left == locatedNode);
        if (locatedNode.right == null && locatedNode.left == null) {
            if (left) {
                locatedNode.father.left = null;
            } else {
                locatedNode.father.right = null;
            } 
        } else if (locatedNode.left != null && locatedNode.right == null) {
            locatedNode.left.father = locatedNode.father;
            if (left) {
                locatedNode.father.left = locatedNode.left;
            } else {
                locatedNode.father.right = locatedNode.left;
            }
        } else if (locatedNode.left == null && locatedNode.right != null) {
            locatedNode.right.father = locatedNode.father;
            if (left) {
                locatedNode.father.left = locatedNode.right;
            } else {
                locatedNode.father.right = locatedNode.right;
            }
        } else if (locatedNode.left != null && locatedNode.right != null) {
            Node min = locatedNode.right;
            //min element has no left or right child.
            if (min.left == null && min.right == null) {
                locatedNode.key = min.key;
                locatedNode.right = null;
                return;
            }
            while (min.left != null) {
                //Min element has left child.
                min = min.left;
            }
            locatedNode.key = min.key;
            if (min.right != null) {
                min.right.father = min.father;
            }
            //Min element has no left child but has right child.
            if (min.father != locatedNode) {
                min.father.left = min.right;
            } else {
                min.father.right = min.right;
            }
            updateHeights(min.father);
            balance(min.father); 
            return;
        }
        updateHeights(locatedNode.father);
        balance(locatedNode.father);
    }

    public void deleteTree() {
        root = null;
    }

    public void printTree(Node node) {
        if (node == null) {
            return;
        }
        printTree(node.left);
        System.out.print(node + "\n");
        printTree(node.right);
    }

    public void consoleUI() {
        Scanner scan = new Scanner(System.in);
        while (true) {
            System.out.println("\n1.- Add items\n"
                    + "2.- Delete items\n"
                    + "3.- Check items\n"
                    + "4.- Print tree\n"
                    + "5.- Delete tree\n");
            int choice = scan.nextInt();

            int item;
            Node node;
            switch (choice) {
                case 1:
                    item = scan.nextInt();
                    while (item != -999) {
                        node = new Node(item, null, null, null);
                        insert(node, root);
                        item = scan.nextInt();
                    }
                    printTree(root);
                    break;
                case 2:
                    item = scan.nextInt();
                    while (item != -999) {
                        node = new Node(item, null, null, null);
                        System.out.print("\nDeleting item " + item);
                        if (deleteNode(node, root)) {
                            System.out.print(": deleted!");
                        } else {
                            System.out.print(": does not exist!");
                        }
                        item = scan.nextInt();
                    }
                    System.out.println();
                    printTree(root);
                    break;
                case 3:
                    item = scan.nextInt();
                    while (item != -999) {
                        node = new Node(item, null, null, null);
                        System.out.println((findNode(node, root) != null) ? "found" : "not found");
                        item = scan.nextInt();
                    }
                    break;
                case 4:
                    printTree(root);
                    break;
                case 5:
                    deleteTree();
                    System.out.println("Tree deleted!");
                    break;
            }
        }
    }
    
    /*AVL Specific Code*/
    private void updateHeights(Node node) {
        if (node == null) {
            return;
        }
        int left = -1, right = -1;
        if (node.left != null) {
            left = node.left.height;
        }
        if (node.right != null) {
            right = node.right.height;
        }
        node.height = Integer.max(left, right) + 1;
        updateHeights(node.father);
    }

    private int checkBalance(Node node){
        int subleft = -1, subright = -1;
        if(node.right!=null)subright = node.right.height;
        if(node.left!=null)subleft = node.left.height;
        int balance = subleft-subright;
        return balance;
    }
    
    private void balance(Node node) {
        int balance = checkBalance(node);
        if (balance < -1) {   
            if(checkBalance(node.right)==1) rotateRight(node.right);
            rotateLeft(node);
        } else if (balance > 1) { 
            if(checkBalance(node.left)==-1) rotateLeft(node.left);
            rotateRight(node);
        } 
        if(node.father!=null) balance(node.father); 
    } 

    void rotateLeft(Node node) {
        //Right heavy - Rotate left 
        System.out.println("Need to rotate node " + node + " left!");
        
        //In case it's a Left-Right rotation
        Node keepAlive = null;
        if(node.left!=null) {keepAlive = node.left; System.out.println("Need to keep alive "+keepAlive.key) ; }
        
        Node newNode = new Node(node.key, null, node.right.left, node);
        node.left = newNode;
        node.key = node.right.key;
        //Handling two special cases
        if (node.right.left != null) node.right.left.father = newNode;
        if (node.right.right != null) node.right.right.father = node; 
        
        node.right = node.right.right;
        if(keepAlive!=null){ node.left.left = keepAlive; keepAlive.father = node.left;}
        updateHeights(newNode);
    }

    void rotateRight(Node node) {
        //Left heavy - Rotate right
        System.out.println("Need to rotate node " + node + " right!");
        
        //In case it's a Right-Left rotation
        Node keepAlive = null;
        if(node.right!= null){ keepAlive = node.right; System.out.println("Need to keep alive "+keepAlive.key);}
        
        Node newNode = new Node(node.key, node.left.right, null, node);
        node.right = newNode;
        node.key = node.left.key;
        //Handle two special cases
        if (node.left.right != null) node.left.right.father = newNode;
        if (node.left.left != null) node.left.left.father = node;
        
        node.left = node.left.left;
        if(keepAlive!=null){ node.right.right = keepAlive; keepAlive.father = node.right;}
        updateHeights(newNode);
    } 
    
    public static void main(String[] args) {
        AVL avl = new AVL();
        avl.consoleUI(); 
    }
}

You can test building and modifying some random AVLs. Here's a link: AVL Example
I've tested it on the AVLs given in the link above. Output shown is on the trees shown in the link.

Output [Note: -999 is the end marker]:

1.- Add items
2.- Delete items
3.- Check items
4.- Print tree
5.- Delete tree

1
15 20 24 10 13 7 30 36 25 -999
Need to rotate node key = 15 height = 2 left!
Need to rotate node key = 10 height = 1 left!
Need to rotate node key = 15 height = 2 right!
Need to rotate node key = 20 height = 3 right!
Need to keep alive 24
Need to rotate node key = 24 height = 2 left!
Need to rotate node key = 30 height = 2 right!
Need to keep alive 36
Need to rotate node key = 20 height = 3 left!
Need to keep alive 15
key = 7 height = 0
key = 10 height = 1
key = 13 height = 3
key = 15 height = 0
key = 20 height = 1
key = 24 height = 2
key = 25 height = 0
key = 30 height = 1
key = 36 height = 0

1.- Add items
2.- Delete items
3.- Check items
4.- Print tree
5.- Delete tree

2
24 -999

Deleting item 24: deleted!
key = 7 height = 0
key = 10 height = 1
key = 13 height = 3
key = 15 height = 0
key = 20 height = 1
key = 25 height = 2
key = 30 height = 1
key = 36 height = 0

1.- Add items
2.- Delete items
3.- Check items
4.- Print tree
5.- Delete tree

2
20 -999

Deleting item 20: deleted!
key = 7 height = 0
key = 10 height = 1
key = 13 height = 3
key = 15 height = 0
key = 25 height = 2
key = 30 height = 1
key = 36 height = 0

1.- Add items
2.- Delete items
3.- Check items
4.- Print tree
5.- Delete tree