数据结构之AVL树

在上一篇博客中，向大家介绍了二叉搜索树这一种数据结构。下面，我们来看一个问题。假设现在有一棵空树，我们先向树中添加一个根节点5，然后再添加6，然后再按顺序添加7、8。我们来看下生成的二叉搜索树是什么样子。

从上面的图中我们可以看出，二叉搜索树退化成了链表，大大的降低了查询效率。为了解决这个问题，我们引入一种新的数据结构——AVL树。

概述

AVL树是一种带有平衡条件的二叉搜索树。它的左右子树的高度差最多为1。如下图所示，左边的是AVL树，而右边的不是。

既然AVL树是一颗自平衡的二叉树搜索树，那么我们在二叉搜索树的基础上维护平衡即可。试想一下，什么情况下会破坏平衡条件呢？

• 插入操作时，向一颗高度差为1的二叉搜索树添加节点，高度差有可能会变成2，破坏了平衡条件。
• 删除操作时，向一颗高度差为1的二叉搜索树删除节点，高度差有可能会变成2，破坏了平衡条件。

通常，我们可以通过在添加节点和删除节点做一些操作来位置平衡，这种操作我们称之为旋转。

旋转

如果平衡条件破坏在树的外侧（左-左、右-右），那么可以通过一次单旋转来修复平衡。如果平衡条件破坏在树的内存（左-右、右-左），那么需要通过一次双旋转来修复平衡。

• “左-左”：可以通过一次右旋转来进行修复平衡。
  

• “右-右”：可以通过一次左旋转来修复平衡。
  

• “右-左”：需要先将Y节点左旋，形成“左-左”的情况，然后再通过一次右旋来修复平衡。
  

• “左-右”：需要先将Y节点右旋，先修复成“右-右”的情况，然后再通过一次左旋来修复平衡。
  

代码

public class AVLTree<K extends Comparable<K>, V> {

    //元素个数
    private int size;

    //跟节点
    private Node root;

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

    private class Node {
        private K key;
        private V value;
        private int height;
        private Node leftChild, rightChild;

        private Node(K key, V value) {
            this(key, value, null, null);
        }

        private Node(K key, V value, Node leftChild, Node rightChild) {
            this.key = key;
            this.value = value;
            this.leftChild = leftChild;
            this.rightChild = rightChild;
            this.height = 1;
        }
    }

    /**
     * 树是否为空
     *
     * @return
     */
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * 树中元素的个数
     *
     * @return
     */
    public int getSize() {
        return size;
    }

    /**
     * 是否包含node.key = key的节点
     *
     * @param key 查找的key
     * @return
     */
    public boolean contains(K key) {
        checkKey(key);
        return find(root, key) == null;
    }

    /**
     * 查找node.key = key的节点
     *
     * @param key
     * @return
     */
    public V get(K key) {
        checkKey(key);
        Node node = find(root, key);
        return node == null ? null : node.value;
    }

    /**
     * 向树中添加元素
     *
     * @param key
     * @param value
     */
    public void add(K key, V value) {
        checkKey(key);
        root = add(root, key, value);
    }

    /**
     * 在以node为根节点的树中添加节点
     *
     * @param node  递归遍历的根节点
     * @param key   新节点的key
     * @param value 新节点的value
     * @return 添加成功，新树的根节点
     */
    private Node add(Node node, K key, V value) {
        if (node == null) {
            size++;
            return new Node(key, value);
        }
        if (key.compareTo(node.key) < 0) {
            //在左子树中添加
            node.leftChild = add(node.leftChild, key, value);
        } else if (key.compareTo(node.key) > 0) {
            //在右子树中添加
            node.rightChild = add(node.rightChild, key, value);
        } else {
            //存在则替换
            node.value = value;
        }
        //节点的高度 = 左右子树的最大值 + 1
        node.height = Math.max(getHeight(node.leftChild), getHeight(node.rightChild)) + 1;
        return node;
    }

    /**
     * 删除元素
     *
     * @param key 元素的key
     * @return
     */
    public V remove(K key) {
        checkKey(key);
        Node retNode = find(root, key);
        if (retNode != null) {
            root = remove(root, key);
        }
        return retNode == null ? null : retNode.value;
    }

    /**
     * 在以node为根节点的树中删除node.key = key的节点
     *
     * @param node 根节点
     * @param key  删除的key
     * @return 删除后，新树的根节点
     */
    private Node remove(Node node, K key) {
        if (node == null) {
            return null;
        }
        Node retNode = null;
        if (key.compareTo(node.key) < 0) {
            //在左子树中删除
            node.leftChild = remove(node.leftChild, key);
            retNode = node;
        } else if (key.compareTo(node.key) > 0) {
            //在右子树中删除
            node.rightChild = remove(node.rightChild, key);
        } else {
            //查找到删除的节点
            if (node.leftChild == null) {
                Node rightChild = node.rightChild;
                node.rightChild = null;
                retNode = rightChild;
                size--;
            } else if (node.rightChild == null) {
                Node leftChild = node.leftChild;
                node.leftChild = null;
                retNode = leftChild;
                size--;
            } else {
                Node min = findMin(node);
                min.rightChild = remove(node.rightChild, min.key);
                min.leftChild = node.leftChild;
                node.leftChild = node.rightChild = null;
                retNode = min;
                size--;
            }
        }
        if (retNode == null) {
            return null;
        }
        retNode.height = Math.max(getHeight(retNode.leftChild), getHeight(retNode.rightChild)) + 1;
        return balance(retNode);
    }


    /**
     * 将以node为根节点的树构建成平衡二叉树
     *
     * @param node 根节点
     * @return 新的平衡二叉树的根节点
     */
    private Node balance(Node node) {
        int balanceFactor = getBalanceFactor(node);
        if (balanceFactor > 1 && getBalanceFactor(node.leftChild) > 0) {
            //LL
            return rightRotate(node);
        } else if (balanceFactor > 1 && getBalanceFactor(node.rightChild) > 0) {
            //LR 先让左子节点进行左旋，变成LL，然后再右旋
            node.leftChild = leftRotate(node.leftChild);
            return rightRotate(node);
        } else if (balanceFactor < -1 && getBalanceFactor(node.leftChild) > 0) {
            //RL 先让右子节点进行右旋，变成RR，然后再左旋
            node.rightChild = rightRotate(node.rightChild);
            return leftRotate(node);
        } else if (balanceFactor < -1 && getBalanceFactor(node.rightChild) > 0) {
            //RR
            return leftRotate(node);
        }
        return node;
    }


    /**
     * 左旋
     *
     * @param x 旋转的根节点
     * @return 旋转后新的根节点
     */
    private Node leftRotate(Node x) {

        //     x                          y
        //	  / \                       /   \
        //   t1  y        左旋         x     z
        //      / \    ---------->    / \     \
        //     t2  z                 t1 t2     t3
        //  	     \
        //          t3

        Node y = x.rightChild;
        Node t2 = y.leftChild;

        x.rightChild = t2;
        y.leftChild = x;

        //因节点发生变化，重新结算高度
        x.height = Math.max(getHeight(x.leftChild), getHeight(x.rightChild));
        y.height = Math.max(getHeight(y.leftChild), getHeight(y.rightChild));

        return y;
    }

    /**
     * 右旋
     *
     * @param x 旋转的根节点
     * @return 旋转后新的根节点
     */
    private Node rightRotate(Node x) {

        //        x                         y
        //       / \                      /   \
        //      y  t1       右旋         z     x
        //     / \       ---------->    /     / \
        //    z  t2                    t3    t2 t1
        //   / 
        //  t3

        Node y = x.leftChild;
        Node t2 = y.rightChild;

        x.leftChild = t2;
        y.rightChild = x;

        x.height = Math.max(getHeight(x.leftChild), getHeight(x.rightChild));
        y.height = Math.max(getHeight(y.leftChild), getHeight(y.rightChild));

        return y;
    }

    /**
     * 节点的高度
     *
     * @param node
     * @return
     */
    private int getHeight(Node node) {
        if (node == null) {
            return 0;
        }
        return node.height;
    }

    /**
     * 计算节点的平衡因子
     *
     * @param node
     * @return
     */
    private int getBalanceFactor(Node node) {
        return getHeight(node.leftChild) - getHeight(node.rightChild);
    }

    /**
     * 在以node为根节点的树中查找node.key = key的节点
     *
     * @param node 根节点
     * @param key  查找的key
     * @return
     */
    private Node find(Node node, K key) {
        if (node == null) {
            //未找到
            return null;
        }
        if (key.compareTo(node.key) < 0) {
            //在左子树中查找
            return find(node.leftChild, key);
        } else if (key.compareTo(node.key) > 0) {
            //在右子树中查找
            return find(node.rightChild, key);
        } else {
            return node;
        }
    }

    /**
     * 在以node节点为根节点的树中查找key最小的节点
     *
     * @param node
     * @return
     */
    private Node findMin(Node node) {
        if (node.leftChild == null) {
            return node;
        }
        return findMin(node.leftChild);
    }


    /**
     * 校验
     *
     * @param key
     */
    private void checkKey(K key) {
        if (key == null) {
            throw new IllegalArgumentException("key is illegal.");
        }
    }
}