AVL树

在计算机科学中，AVL树是最早被发明的自平衡二叉查找树。在AVL树中，任一节点对应的两棵子树的最大高度差为1，因此它也被称为高度平衡树。查找、插入和删除在平均和最坏情况下的时间复杂度都是{\displaystyle O(\log {n})}O(\log{n})。增加和删除元素的操作则可能需要借由一次或多次树旋转，以实现树的重新平衡。AVL树得名于它的发明者G. M. Adelson-Velsky和Evgenii Landis，他们在1962年的论文《An algorithm for the organization of information》中公开了这一数据结构。

节点的平衡因子是它的左子树的高度减去它的右子树的高度（有时相反）。带有平衡因子1、0或 -1的节点被认为是平衡的。带有平衡因子 -2或2的节点被认为是不平衡的，并需要重新平衡这个树。平衡因子可以直接存储在每个节点中，或从可能存储在节点中的子树高度计算出来。

import java.io.Serializable;

public class AVLTree<K extends Comparable<K> & Serializable, V> {
    private class Node {
        K key;
        V value;
        int height;
        Node left, right;

        Node(K key, V value) {
            this.key = key;
            this.value = value;
            this.height = 1;
            left = right = null;
        }
    }

    private Node root;
    private int size;

    public AVLTree() {
        root = null;
        size = 0;
    }

    public boolean isEmpty() {
        return size == 0;
    }

    public int size() {
        return size;
    }

    public void add(K key, V value) {
        root = add(root, key, value);
    }

    private Node add(Node node, K key, V value) {
        if (node == null) {
            size++;
            return new Node(key, value);
        }

        if (node.key.compareTo(key) > 0) {
            node.left = add(node.left, key, value);
        } else if (node.key.compareTo(key) < 0) {
            node.right = add(node.right, key, value);
        } else {
            node.value = value;
        }
        node.height = Math.max(getHeight(node.left), getHeight(node.right)) + 1;

        int balanceFactor = getBalanceFactor(node);
        if (balanceFactor > 1 && getBalanceFactor(node.left) >= 0) {
            return rightRotate(node);
        }

        if (balanceFactor < -1 && getBalanceFactor(node.right) <= 0) {
            return leftRotate(node);
        }

        if (balanceFactor > 1 && getBalanceFactor(node.left) < 0) {
            node.left = leftRotate(node.left);
            return rightRotate(node);
        }

        if (balanceFactor < -1 && getBalanceFactor(node.right) > 0) {
            node.right = rightRotate(node.right);
            return leftRotate(node);
        }

        return node;
    }

    public V remove(K key) {
        Node node = getNode(key);
        if (node != null) {
            root = remove(root, key);
            return node.value;
        }
        return null;
    }

    private Node getNode(K key) {
        return getNode(root, key);
    }

    private Node getNode(Node node, K key) {
        if (node == null) {
            return null;
        }
        if (node.key.compareTo(key) == 0) {
            return node;
        } else if (node.key.compareTo(key) > 0) {
            return getNode(node.left, key);
        } else {
            return getNode(node.right, key);
        }
    }

    private Node remove(Node node, K key) {
        if (node == null) {
            return null;
        }
        Node ret;
        if (node.key.compareTo(key) > 0) {
            node.left = remove(node.left, key);
            ret = node;
        } else if (node.key.compareTo(key) < 0) {
            node.right = remove(node.right, key);
            ret = node;
        } else {
            if (node.left == null) {
                Node right = node.right;
                node.right = null;
                size--;
                ret = right;
            } else if (node.right == null) {
                Node left = node.left;
                node.left = null;
                size--;
                ret = left;
            } else {
                Node successor = minimum(node.right);
                successor.left = node.left;
                successor.right = removeMin(node.right);
                node.left = node.right = null;
                ret = successor;
            }
        }
        if (ret == null) {
            return null;
        }

        ret.height = Math.max(getHeight(ret.left), getHeight(ret.right)) + 1;

        int balanceFactor = getBalanceFactor(ret);

        if (balanceFactor > 1 && getBalanceFactor(ret.left) >= 0) {
            return rightRotate(ret);
        }

        if (balanceFactor < -1 && getBalanceFactor(ret.right) <= 0) {
            return leftRotate(ret);
        }

        if (balanceFactor > 1 && getBalanceFactor(ret.left) < 0) {
            ret.left = leftRotate(ret.left);
            return rightRotate(ret);
        }

        if (balanceFactor < -1 && getBalanceFactor(ret.right) > 0) {
            ret.right = rightRotate(ret.right);
            return leftRotate(ret);
        }

        return ret;
    }

    private Node removeMin(Node node) {
        if (node.left == null) {
            Node right = node.right;
            node.right = null;
            size--;
            return right;
        }
        node.left = removeMin(node.left);
        return node;
    }

    private Node minimum(Node node) {
        if (node.left == null) {
            return node;
        }
        return minimum(node.left);
    }

    private Node leftRotate(Node a) {
        Node b = a.right;
        a.right = b.left;
        b.left = a;

        a.height = Math.max(getHeight(a.left), getHeight(a.right)) + 1;
        b.height = Math.max(getHeight(b.left), getHeight(b.right)) + 1;
        return b;
    }

    private Node rightRotate(Node a) {
        Node b = a.left;
        a.left = b.right;
        b.right = a;

        a.height = Math.max(getHeight(a.left), getHeight(a.right)) + 1;
        b.height = Math.max(getHeight(b.left), getHeight(b.right)) + 1;
        return b;
    }

    private int getBalanceFactor(Node node) {
        if (node == null) {
            return 0;
        }
        return getHeight(node.left) - getHeight(node.right);
    }

    private int getHeight(Node node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    private boolean isBalanced() {
        return isBalanced(root);
    }

    private boolean isBalanced(Node node) {
        if (node == null) {
            return true;
        }
        int balanceFactor = getBalanceFactor(node);
        if (Math.abs(balanceFactor) > 1) {
            return false;
        }
        return isBalanced(node.left) && isBalanced(node.right);
    }

    public static void main(String[] args) {
        AVLTree<Integer, Integer> avl = new AVLTree<>();
        for (int i = 0; i < 100; i++) {
            avl.add(i, i);
        }
        System.out.println(avl.isBalanced());

        System.out.println(avl);
        System.out.println(avl.size());

        for (int i = 0; i < 100; i++) {
            avl.remove(i);
        }

        System.out.println(avl.size());

        System.out.println(avl);
    }

    @Override
    public String toString() {
        StringBuilder buf = new StringBuilder();
        generate(buf, root);
        return buf.toString();
    }

    private void generate(StringBuilder buf, Node node) {
        if (node == null) {
            return;
        }
        generate(buf, node.left);
        buf.append(node.value).append(" ");
        generate(buf, node.right);
    }
}