AVL树，self

package com.company.stucts.binarytree.searchtree.balancetree;/**来自 N o w J a v a . c o m**/


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;

    }


    public void add(K key, V value) {

        Node<K, V> node = rootNode;

        Node<K, V> parentNode = null;/** 来 自 NowJava.com - 时代Java**/


        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);

        }