AVL tree, self-balancing binary search tree.

Description

Demo Code

package com.company.stucts.binarytree.searchtree.balancetree;

import java.util.Comparator;

@SuppressWarnings({ "Duplicates", "WeakerAccess" })
public class AVLTree<K, V> {
    private Node<K, V> rootNode;
    private final Comparator<K> comparator;

    public AVLTree(Comparator<K> comparator) {
        this.comparator = comparator;
    }/*w ww .jav a 2  s.  c  om*/

    public void add(K key, V value) {
        Node<K, V> node = rootNode;
        Node<K, V> parentNode = null;

        while (node != null) {
            int result = comparator.compare(key, node.key);
            if (result == 0) {
                node.value = value;
                return;
            } else {
                parentNode = node;
                if (result < 0) {
                    node = node.left;
                } else {
                    node = node.right;
                }
            }
        }

        Node<K, V> newNode = new Node<>(key, value);
        newNode.parent = parentNode;

        if (parentNode == null) {
            rootNode = newNode;
        } else {
            if (comparator.compare(key, parentNode.key) < 0) {
                parentNode.left = newNode;
            } else {
                parentNode.right = newNode;
            }
        }

        // fix tree balance if needed
        updateBalance(newNode);
    }

    private void balanceTree(Node<K, V> imbalancedNode) {
        System.out.println("Imbalance node: [" + imbalancedNode.key + ","
                + imbalancedNode.value + "]" + ", balance: "
                + imbalancedNode.balance);
        if (imbalancedNode.balance == -2) {
            if (imbalancedNode.left.balance == 1) {
                System.out.println("New node in left subtree, right child");
                System.out.println("Perform left-right rotation");
                rotateLeft(imbalancedNode.right);
                rotateRight(imbalancedNode);
            } else {
                System.out.println("New node in left subtree, left child");
                System.out.println("Perform right rotation");
                rotateRight(imbalancedNode);
            }
        } else {
            if (imbalancedNode.right.balance == -1) {
                System.out.println("New node in right subtree, left child");
                System.out.println("Perform right-left rotation");
                rotateRight(imbalancedNode.right);
                rotateLeft(imbalancedNode);
            } else {
                System.out
                        .println("New node in right subtree, right child");
                System.out.println("Perform left rotation");
                rotateLeft(imbalancedNode);
            }
        }
    }

    private void rotateRight(Node<K, V> imbalanceNode) {
        Node<K, V> leftNode = imbalanceNode.left;
        imbalanceNode.left = leftNode.right;
        if (leftNode.right != null) {
            leftNode.right.parent = imbalanceNode;
        }

        leftNode.parent = imbalanceNode.parent;
        if (imbalanceNode.parent == null) {
            rootNode = leftNode;
        } else if (imbalanceNode.parent.right == imbalanceNode) {
            imbalanceNode.parent.right = leftNode;
        } else {
            imbalanceNode.parent.left = leftNode;
        }

        leftNode.right = imbalanceNode;
        imbalanceNode.parent = leftNode;

        // newBal(B) = oldBal(B) + 1 - min(0, oldBal(D)
        imbalanceNode.balance = imbalanceNode.balance + 1
                - Math.min(0, leftNode.balance);
        // newBal(D) = oldBal(D) + 1 + max(0, newBal(B)
        leftNode.balance = leftNode.balance + 1
                + Math.max(0, imbalanceNode.balance);
    }

    private void rotateLeft(Node<K, V> imbalanceNode) {
        Node<K, V> rightNode = imbalanceNode.right;
        imbalanceNode.right = rightNode.left;
        if (rightNode.left != null) {
            rightNode.left.parent = imbalanceNode;
        }

        rightNode.parent = imbalanceNode.parent;
        if (imbalanceNode.parent == null) {
            rootNode = rightNode;
        } else if (imbalanceNode.parent.left == imbalanceNode) {
            imbalanceNode.parent.left = rightNode;
        } else {
            imbalanceNode.parent.right = rightNode;
        }

        rightNode.left = imbalanceNode;
        imbalanceNode.parent = rightNode;

        // newBal(B) = oldBal(B) - 1 - max(0, oldBal(D)
        imbalanceNode.balance = imbalanceNode.balance - 1
                - Math.max(0, rightNode.balance);
        // newBal(D) = oldBal(D) - 1 + min(0, newBal(B)
        rightNode.balance = rightNode.balance - 1
                + Math.min(0, imbalanceNode.balance);
    }

    private void updateBalance(Node<K, V> node) {
        while (node.parent != null) {
            if (node.parent.left == node) {
                node.parent.balance -= 1;
            } else {
                node.parent.balance += 1;
            }

            if (Math.abs(node.parent.balance) > 1) {
                balanceTree(node.parent);
                return;
            } else if (node.parent.balance == 0) {
                return;
            } else {
                node = node.parent;
            }
        }
    }

    public void printTree() {
        printNode(rootNode);
    }

    private void printNode(Node<K, V> node) {
        System.out.print("Node: [" + node.key + "," + node.value + "]"
                + ", balance: " + node.balance);
        if (node.parent != null) {
            System.out.println(", -> Parent: [" + node.parent.key + ","
                    + node.parent.value + "]" + ", balance: "
                    + node.parent.balance);
        } else {
            System.out.println(", -> Parent: [ null ]");
        }

        if (node.left != null) {
            printNode(node.left);
        }

        if (node.right != null) {
            printNode(node.right);
        }
    }

    private static class Node<K, V> {
        private K key;
        private V value;
        private int balance;
        private Node<K, V> left, right, parent;

        Node(K key, V value) {
            this.key = key;
            this.value = value;
        }
    }

    /**
     * Test main method
     */
    public static void main(String[] args) {
        Comparator<Integer> comparator = new Comparator<Integer>() {
            @Override
            public int compare(Integer o1, Integer o2) {
                return o1 < o2 ? -1 : (o1 > o2 ? 1 : 0);
            }
        };

        AVLTree<Integer, Integer> tree = new AVLTree<>(comparator);
        tree.add(1, null);
        tree.add(2, null);
        tree.add(3, null);
        tree.add(4, null);
        tree.printTree();
    }
}
AVL tree, self-balancing binary search tree. - Java Data Structure

Description

Demo Code

Related Tutorials
