toys/data-structures/avl-tree/AVLTree.java

/**
 * Made up out of thin air, because all the resources for this are crap.
 */

public class AVLTree<T> extends LinkedBinarySearchTree<T> {
    /**
     * Constructors
     */
    public AVLTree () {
        super ();
    }
    public AVLTree (T key, T element) {
        super (key, element);
    }
    /**
     * insert
     * Calls balance()
     */
    public void insert (T key, T element) {
        Node<T> temp = new Node<T> (key, element);
        this.addElement (temp);
        
        /*
         * Add balance factors
         */
        Node<T> current = temp;
        while (current != root) {
            if (current.isRightChild()) {
                if ((current.getParent().getHeight() == 0)) {
                    current.getParent().setRightHeight(1);
                } else {// This is an internal node
                    current.getParent().setRightHeight(current.getHeight() + 1);
                }
            } else if (current.isLeftChild()) {
                if ((current.getParent().getHeight() == 0)) {
                    current.getParent().setLeftHeight(1);
                } else {// This is an internal node
                    current.getParent().setLeftHeight(current.getHeight() + 1);
                }
            }
            current = current.getParent();
        }
        
        this.balance(temp);
    }
    /**
     * remove
     * Calls balance()
     */
    public void remove (Node<T> element) {
        Node<T> current = element.getParent();
        if (element.isLeaf()) {
            if (element == root)
                this.removeElement (element.getKey());
            else {
                if (element.isLeftChild()) {
                    this.removeElement (element.getKey());
                    current.setLeftHeight(0);
                } else {
                    this.removeElement (element.getKey());
                    current.setRightHeight(0);
                }
            }
        } else {
            current = element;
            Node<T> parent = element.getParent();
            
            /*
             * Get the correct replacement
             */
            if (current.hasRight()) {
                current = current.getRight();
                if (current.hasLeft()) {
                    while (current.hasLeft())
                        current = current.getLeft();
                    if (current.hasRight())
                        current.getParent().setLeftHeight(current.getRight().getHeight() + 1);
                    else
                        current.getParent().setLeftHeight(0);
                    current.getParent().setLeft(current.getRight());
                }
            } else {
                current = current.getRight();
                if (current.hasRight()) {
                    while (current.hasRight())
                        current = current.getRight();
                    if (current.hasLeft())
                        current.getParent().setRightHeight(current.getLeft().getHeight() + 1);
                    else
                        current.getParent().setRightHeight(0);
                    current.getParent().setRight(current.getLeft());
                }
            }
            
            /*
             * Set the replacement's numbers and relationships
             */
            if (element.getRight() != current)
                current.setRight(element.getRight());
            if (element.getLeft() != current)
                current.setLeft(element.getLeft());
            if (current.getRight() != null) {
                current.getRight().setParent(current);
                current.setRightHeight(current.getRight().getHeight() + 1);
            } else
                current.setRightHeight(0);
            if (current.getLeft() != null) {
                current.getLeft().setParent(current);
                current.setLeftHeight(current.getLeft().getHeight() + 1);
            } else
                current.setLeftHeight(0);
            if (element.isRightChild())
                parent.setRight(current);
            else if (element.isLeftChild())
                parent.setLeft(current);
            current.setParent(parent);
            
            /*
             * Delete the element
             */
            this.removeElement (element.getKey());
        }
        if (!isEmpty())
            this.balance(current);
    }
    /**
     * balance
     * Calls rotateRight() and rotateLeft()
     */
    private void balance (Node<T> current) {
        do {
            if (Math.abs(current.getBalance()) > 1) {
                if (current.getBalance() > 1) {
                    // Right subtree unbalanced
                    if (current.getRight().getBalance() == 1) {
                        // Problem is in the right child's right subtree
                        this.rotateLeft(current);
                    } else {
                        // Problem is in the right child's left subtree
                        this.rotateRight(current.getRight());
                        this.rotateLeft(current);
                    }
                } else {
                    // Left subtree unbalanced
                    if (current.getLeft().getBalance() == 1) {
                        // Problem is in the left child's right subtree
                        this.rotateLeft(current.getLeft());
                        this.rotateRight(current);
                    } else {
                        // Problem is in the left child's left subtree
                        this.rotateRight(current);
                    }
                }
            }
            current = current.getParent();
        } while (current != null);
    }
    /**
     * rotateRight
     */
    private void rotateRight(Node<T> temp) {
        Node<T> current = temp.getLeft();
        Node<T> parent  = temp;
        
        /*
         * Problem is in the left child's left subtree
         */
        
        // Step 1: Change current's parent
        if (parent.getParent() == null) {
            current.setParent(null);
            current.unsetIsRight();
            current.unsetIsLeft();
            root = current;
        } else {
            if (parent.isRightChild())
                parent.getParent().setRight(current);
            else if (parent.isLeftChild())
                parent.getParent().setLeft(current);
        }
        
        // Step 2: Change parent's left
        parent.setLeft(current.getRight());
        if (current.getRight() != null)
            parent.getLeft().setParent(parent);
        
        // Step 3: Change current's right
        current.setRight(parent);
        parent.setParent(current);
        
        /*
         * Change balance factors
         */
        if (parent.getLeft() != null)
            parent.setLeftHeight(parent.getLeft().getHeight() + 1);
        else
            parent.setLeftHeight(0);
        current.setRightHeight(parent.getHeight() + 1);
        if (current.getParent() != null) {
            if (current.isRightChild())
                current.getParent().setRightHeight(current.getHeight() + 1);
            else if (current.isLeftChild())
                current.getParent().setLeftHeight(current.getHeight() + 1);
        }
    }
    /**
     * rotateLeft
     */
    private void rotateLeft(Node<T> temp) {
        Node<T> current = temp.getRight();
        Node<T> parent  = temp;
        
        /*
         * Problem is in right child's right subtree
         */
        
        // Step 1: Change current's parent
        if (parent.getParent() == null) {
            current.setParent(null);
            current.unsetIsRight();
            current.unsetIsLeft();
            root = current;
        } else {
            if (parent.isRightChild())
                parent.getParent().setRight(current);
            else if (parent.isLeftChild())
                parent.getParent().setLeft(current);
            current.setParent(parent.getParent());
        }
        
        // Step 2: Change parent's right
        parent.setRight(current.getLeft());
        if (current.getLeft() != null)
            parent.getRight().setParent(parent);
        
        // Step 3: Change current's left
        current.setLeft(parent);
        parent.setParent(current);
        
        /*
         * Change balance factors
         */
        if (parent.getRight() != null)
            parent.setRightHeight(parent.getRight().getHeight() + 1);
        else
            parent.setRightHeight(0);
        current.setLeftHeight(parent.getHeight() + 1);
        if (current.getParent() != null) {
            if (current.isRightChild())
                current.getParent().setRightHeight(current.getHeight() + 1);
            else if (current.isLeftChild())
                current.getParent().setLeftHeight(current.getHeight() + 1);
        }
    }
}