数据结构中的Boss - AVL平衡二叉树

前言

在讨论平衡二叉树之前，先了解什么是二叉搜索树，二叉搜索树是二叉树的一种，它有一种特性，就是每个节点的左子节点都比节点本身的值小，右子节点都比节点本身大。因为这个特性，当对二叉搜索树进行中序遍历的时候，输出一定是按升序排列的。
利用二叉搜索树，可以在$ O(log N) $的时间复杂度下查找指定元素。然而如果在插入二叉搜索树的时候，是以升序的方式，比如[1,2,3,4,5]，二叉搜索树会一直往右节点增加，这样会导致二叉搜索树退化成链表，查找的时间复杂度也降到了$ O(N) $。

为了保持二叉搜索树的特性，同时让树的高度尽量缩小，于是平衡二叉搜索树（简称：平衡二叉树）产生了。所谓平衡，就是让树的左右子节点的高度尽量相等，左右两边尽量保持平衡，使节点平均分布在两侧，这样使得查找效率始终维持在$ O(log N) $，也就是树的高度。

定义

平衡二叉搜索树是由前苏联的数学家 Adelse-Velskil 和 Landis 在 1962 年最先提出的高度平衡的二叉树，根据科学家的英文名所以也被称为 AVL树。

百度百科：在计算机科学中，AVL树是最先发明的自平衡二叉查找树。在AVL树中任何节点的两个子树的高度最大差别为1，所以它也被称为高度平衡树。增加和删除可能需要通过一次或多次树旋转来重新平衡这个树。AVL树得名于它的发明者G. M. Adelson-Velsky和E. M. Landis，他们在1962年的论文《An algorithm for the organization of information》中发表了它。

它具有如下几个性质：

本身首先是一棵二叉搜索树。
带有平衡条件：每个节点的左右子树的高度之差（平衡因子）的绝对值最多为1。

也就是说，AVL树本质上是带了平衡功能的二叉搜索树。除此之外，还有其他实现了平衡搜索树的数据结构，比如：B树，树堆(Treap)，以及鼎鼎有名的红黑树等。

平衡因子

我们规定平衡因子为节点的左右子树的高度之差，即左节点高度-右节点高度，英文名称Balance Factor。通过这个我们不仅可以判断树是否平衡，而且可以根据符号的正负判断是左子树高还是右子树高，平衡因子为正表示左子树高，为负表示右子树高，为0表示2边高度相等。

一棵AVL平衡树的所有节点的平衡因子绝对值都不会超过1，下面列举几个例子：

图1：平衡二叉树

图2：不是平衡二叉树

图3：不是平衡二叉树

图4：平衡二叉树

图1是一棵标准的平衡二叉树，它满足二叉搜索树的条件，同时它的每个节点的平衡因子绝对值都不超过1
图2虽然满足平衡因子的条件，但是节点6比节点3大，节点3的所有左子节点都应该比3小，不满足二叉搜索树条件
图3节点5的平衡因子等于3（左节点高度是3，无右节点高度是0），不满足平衡因子条件
图4是平衡二叉树

最小不平衡子树

叶子节点的左右子节点为空，所以它的平衡因子是0。当我们插入一个节点时，树的高度发生变化，距离这个节点最近的，且平衡因子的绝对值大于1的节点，称之为这棵二叉平衡树的最小不平衡子树，其中该节点为子树的根节点。

例如：当我们在这棵平衡二叉树上插入节点7的时候，根据二叉搜索树的规则，每个新插入的节点都会插入到叶子节点，这里会插入到节点9的左节点，根据定义，距离节点7最近的平衡因子绝对值超过1的节点是节点10，节点10就是导致二叉树失衡的最小不平衡子树，同时它是这棵子树的根节点。

引入最小不平衡子树的目的就在于，一旦我们将这棵子树调节平衡，它的父节点、父父节点，直到整棵二叉树的根节点都会变得平衡。因为在不改变有序的前提下，将这棵子树调节平衡，它的高度一定会减少，进而影响它的父节点的平衡因子，而每个节点的平衡因子都是由它的左右子节点高度决定的，所以会自底向上，一步步影响到根节点。

AVL树的关键就是在二叉搜索树的基础上增加了对最小不平衡子树的调整失衡处理。

调整失衡

当插入或者删除节点时，都可能导致树的高度发生变化，引起平衡因子改变使之失衡。为了保持树的平衡，并且维持二叉搜索树的有序特性，AVL树的原理就是对最小不平衡子树进行旋转。就像游戏关卡中的boss一样，每个boss都会有它的弱点，请牢牢记住左旋转和右旋转节点是如何发生变化的，这就是AVL树的弱点，也即核心思想，所有的调整都依赖这2个操作。

我们先定义一个二叉树的节点数据结构：

public class Node
{
    public TKey key;
    public Node left, right;
    public int height { get; private set; }

    public Node(TKey key)
    {
        this.key = key;
        height = 1;
    }

    public int balance_fact()
    {
        var lh = left?.height ?? 0;
        var rh = right?.height ?? 0;
        return lh - rh;
    }

    public void update()
    {
        height = Math.Max(left?.height ?? 0, right?.height ?? 0) + 1;
    }
}

节点包含一个左子节点left和右子节点right，一个记录自身高度的height属性，它自身存储的信息用了泛型TKey，可以是任意类型，只要类型之间可以比较大小（有序的前提当然是可以比较大小）。封装了一个balance_fact方法用于动态获取平衡因子，update方法封装了更新自身高度的代码。

左旋转

当因右节点高度过大导致的失衡（平衡因子 < -1）时，我们进行左旋转，即从最小不平衡子树的根节点(节点3)开始，将根节点的右子节点(节点4)提升为新的根节点，同时它自身成为新根节点的左节点，

当然还存在一种情况，如果右子节点本身就存在左节点，例如下图的节点3，那么它自己的左节点应该放在哪里呢？为了保持有序，那么它的左节点应该成为它原先根节点的右节点，像这样：

转换成代码为：

private Node left_rotate(Node root)
{
    var right = root.right;
    root.right = root.left;
    right.left = root;
    root.update();
    right.update();
    return right;
}

注意这里旋转后，根节点和右子节点的高度都发生了变化，这里需要对高度进行更新操作，由于更新高度是从子节点自底向上更新的，这里根节点变成了右子节点的左节点，所以应先对root节点进行更新，再对right节点更新。

右旋转

当因左节点高度过大导致的失衡（平衡因子 > 1）时，我们需要进行右旋转，右旋转和左旋转是完全镜像的操作，这里不再过多赘述。

private Node right_rotate(Node root)
{
    var left = root.left;
    root.left = left.right;
    left.right = root;
    root.update();
    left.update();
    return left;
}

先左旋后右旋

再来看这样一种情况，如下图所示，当我们在AVL树中插入节点2的时候，节点3变成了最小不平衡子树，由于最小不平衡子树是在左节点，按之前的规则我们应该进行右旋转，旋转之后，我们发现虽然各节点达到平衡，但是节点2却变成了节点1的左节点，已经不满足二叉搜索树左节点必须小于自身节点的性质了。

和之前的直接右旋进行比较可以发现区别，只有当左节点的平衡因子是1时，右旋操作才能保证正确，而如果插入操作是在左节点的右边时，左节点的平衡因子为-1，也即和根节点的平衡因子符号相反，这时并不能直接右旋。

对于这种情况我们需要进行2步旋转操作，即先对左节点左旋，再对根节点右旋。

我们并没有引入额外的程序逻辑，依旧使用最原始的2个旋转操作，只是多了一个旋转步骤而已。

private Node left_right_rotate(Node root)
{
    root.left = left_rotate(root.left);
    return right_rotate(root);
}

先右旋后左旋

同理，在下图情况下，我们需先对右节点右旋，再对根节点左旋。第1次旋转，让平衡因子符号相同，第2次旋转消除失衡。

private Node right_left_rotate(Node node)
{
    node.right = right_rotate(node.right);
    return left_rotate(node);
}

总结：以插入操作为例，可能会产生以下四种失衡情况：

插入方式	描述	旋转方式
LL	在AVL树的左子树根节点的左子树上插入节点而破坏平衡	右旋转
RR	在AVL树的右子树根节点的右子树上插入节点而破坏平衡	左旋转
LR	在AVL树的左子树根节点的右子树上插入节点而破坏平衡	先左旋后右旋
RL	在AVL树的右子树根节点的左子树上插入节点而破坏平衡	先右旋后左旋

对于调整失衡，无非是针对左子树的调整或者对右子树的调整，在调整过程中，再根据左节点或者右节点的平衡因子判断是进行一次旋转或是二次旋转，所以我们再封装2个调整失衡的方法：

private Node left_balance(Node root)
{
    var left = root.left;
    if (left.balance_fact() == -1)
        return left_right_rotate(root);
    return right_rotate(root);
}

private Node right_balance(Node root)
{
    var right = root.right;
    if (right.balance_fact() == 1)
        return right_left_rotate(root);
    return left_rotate(root);
}

当左子树失衡时，我们进行左平衡，左平衡的核心是右旋操作，如果左子节点的平衡因子为负，则先左旋再右旋；当右子树失衡时，我们进行右平衡，右平衡的核心是左旋操作，如果右子节点的平衡因子为正，则先右旋再左旋。前面已经说到，我们是针对最小不平衡子树调整失衡的，所以传递进方法里的参数一定是最小不平衡子树的根节点，而它的左右节点的平衡因子一定只有0、-1、1三种状态，所以这里判断正负用1和-1即可。

插入节点

有了left_balance和right_balance方法，AVL树的插入操作基本和普通的二叉搜索树的插入类似，我们只需要在插入过程中，额外的对插入后的平衡因子进行判断，然后再相应的做左平衡或者右平衡调整即可。

private Node put(Node node, TKey key)
{
    if (node == null)
        return new Node(key);
    int cmp = compare(key, node.key);
    if (cmp < 0)
    {
        node.left = put(node.left, key);
        if (node.balance_fact() > 1)
            node = left_balance(node);
    }
    else if (cmp > 0)
    {
        node.right = put(node.right, key);
        if (node.balance_fact() < -1)
            node = right_balance(node);
    }
    node.update();
    return node;
}

compare方法是针对泛型TKey比较的代码。

删除节点

删除节点的递归调用和插入节点类似，如果比当前节点小，说明要删除的节点在左子树，删除之后可能会引起右子树高度大于左子树高度（平衡因子 < -1），所以要进行右平衡；如果比当前节点大，说明删除的节点在右子树，同理，左子树高度会大于右子树高度，我们判断平衡因子是否大于1来决定是否左平衡。

当值相等时，表示找到要删除的节点，如果这个节点，只包含左节点或者右节点，我们可以将它一边的子节点直接提升为根节点，即返回那个它唯一拥有的子节点，如果不包含子节点，我们就直接返回空，而如果它同时具备2个子节点，这个时候就稍微有点复杂了。AVL树的策略是先找到删除节点右子树的最小节点，即右子树最靠左的那个节点，从右子树中移除这个节点，把它放到根节点的位置，最后返回它。

例如下图中的AVL树要删除节点6，先找到右子树最靠左的节点，也就是右子树最小的节点，这里是节点7，在右子树中删除节点7，然后将节点6的左右节点赋值给节点7。

private Node delete(Node node, TKey key)
{
    if (node == null)
        throw new KeyNotFoundException("给定关键字不在字典中。");
    int cmp = compare(key, node.key);
    if (cmp < 0)
    {
        node.left = delete(node.left, key);
        if (node.balance_fact() < -1)
            node = right_balance(node);
    }
    else if (cmp > 0)
    {
        node.right = delete(node.right, key);
        if (node.balance_fact() > 1)
            node = left_balance(node);
    }
    else
    {
        if (node.right == null)    //右节点为空,直接返回左节点
            return node.left;
        if (node.left == null)     //左节点为空,直接返回右节点
            return node.right;
        Node temp = node;
        node = min(temp.right);
        node.right = deleteMin(temp.right);
        node.left = temp.left;
    }
    node.update();
    return node;
}

这里用到了2个新方法，获取最小节点的min方法和删除最小节点的deleteMin方法。

min方法和普通的二叉搜索树没有区别，就是一直递归查找最左边的节点。我们用这个方法找到节点7

private Node min(Node x)
{
    if (x.left == null)
        return x;
    return min(x.left);
}

deleteMin方法首先判断节点本身是否包含左节点，如果没有，则将它的右节点提升为根节点，相当于它自身被删除了，如果有左节点则传递左节点递归调用，删除完毕之后，左子树高度可能会减小，此时判断平衡因子，看是否需要进行右平衡。我们用这个方法删除右子树的节点7

private Node deleteMin(Node node)
{
    if (node.left == null)
        return node.right;
    node.left = deleteMin(node.left);
    if (node.balance_fact() < -1)
        node = right_balance(node);
    node.update();
    return node;
}

完整代码

因为AVL树本质上就是一棵二叉搜索树，有了关键的插入和删除代码，关于搜索的代码其实就和普通的二查搜索树一样了。

public class AvlBST<TKey>
{
    public class Node
    {
        public TKey key;
        public Node left, right;
        public int height { get; private set; }
        public int size { get; private set; }

        public Node(TKey key)
        {
            this.key = key;
            height = 1;
            size = 1;
        }

        public int balance_fact()
        {
            var lh = left?.height ?? 0;
            var rh = right?.height ?? 0;
            return lh - rh;
        }

        public void update()
        {
            height = Math.Max(left?.height ?? 0, right?.height ?? 0) + 1;
            size = (left?.size ?? 0) + (right?.size ?? 0) + 1;
        }
    }

    /// <summary>
    /// 左旋操作
    /// </summary>
    /// <param name="node"></param>
    /// <returns></returns>
    private Node left_rotate(Node root)
    {
        var right = root.right;
        root.right = root.left;
        right.left = root;
        root.update();
        right.update();
        return right;
    }

    /// <summary>
    /// 右旋操作
    /// </summary>
    /// <param name="node"></param>
    /// <returns></returns>
    private Node right_rotate(Node root)
    {
        var left = root.left;
        root.left = left.right;
        left.right = root;
        root.update();
        left.update();
        return left;
    }

    private Node left_right_rotate(Node root)
    {
        root.left = left_rotate(root.left);
        return right_rotate(root);
    }

    private Node right_left_rotate(Node node)
    {
        node.right = right_rotate(node.right);
        return left_rotate(node);
    }

    private Node left_balance(Node root)
    {
        var left = root.left;
        if (left.balance_fact() == -1)
            return left_right_rotate(root); 
        return right_rotate(root);
    }

    private Node right_balance(Node root)
    {
        var right = root.right;
        if (right.balance_fact() == 1)
            return right_left_rotate(root);
        return left_rotate(root);
    }

    private Node root;
    private Func<TKey, TKey, int> compare;

    public AvlBST(IComparer<TKey> comparer)
    {
        compare = comparer.Compare;
    }

    public void delete(TKey key)
    {
        root = delete(root, key);
    }

    private Node delete(Node node, TKey key)
    {
        if (node == null)
            throw new KeyNotFoundException("给定关键字不在字典中。");
        int cmp = compare(key, node.key);
        if (cmp < 0)
        {
            node.left = delete(node.left, key);
            if (node.balance_fact() < -1)
                node = right_balance(node);
        }
        else if (cmp > 0)
        {
            node.right = delete(node.right, key);
            if (node.balance_fact() > 1)
                node = left_balance(node);
        }
        else
        {
            if (node.right == null)
                return node.left;
            if (node.left == null)
                return node.right;
            Node temp = node;
            node = min(temp.right);
            node.right = deleteMin(temp.right);
            node.left = temp.left;
        }
        node.update();
        return node;
    }

    private Node deleteMin(Node node)
    {
        if (node.left == null)
            return node.right;
        node.left = deleteMin(node.left);
        if (node.balance_fact() < -1)
            node = right_balance(node);
        node.update();
        return node;
    }

    public bool isEmpty()
    {
        return root == null;
    }


    public TKey max()
    {
        if (root == null)
            throw new IndexOutOfRangeException();
        return max(root).key;
    }

    private Node max(Node node)
    {
        if (node.right == null)
            return node;
        return max(node.right);
    }

    public TKey min()
    {
        if (root == null)
            throw new IndexOutOfRangeException();
        return min(root).key;
    }

    private Node min(Node x)
    {
        if (x.left == null)
            return x;
        return min(x.left);
    }

    public void put(TKey key)
    {
        root = put(root, key);
    }

    private Node put(Node node, TKey key)
    {
        if (node == null)
            return new Node(key);
        int cmp = compare(key, node.key);
        if (cmp < 0)
        {
            node.left = put(node.left, key);
            if (node.balance_fact() > 1)
                node = left_balance(node);
        }
        else if (cmp > 0)
        {
            node.right = put(node.right, key);
            if (node.balance_fact() < -1)
                node = right_balance(node);
        }
        node.update();
        return node;
    }

    /// <summary>
    /// 元素key的排位
    /// </summary>
    /// <param name="key"></param>
    /// <returns></returns>
    public int rank(TKey key)
    {
        return rank(root, key);
    }

    private int rank(Node node, TKey key)
    {
        if (node == null)
            return 0;
        int cmp = compare(key, node.key);
        if (cmp == 0)
            return size(node.left);
        if (cmp < 0)
            return rank(node.left, key);
        else
            return size(node.left) + 1 + rank(node.right, key);
    }

    /// <summary>
    /// 查找第k个元素
    /// </summary>
    /// <param name="k">k从0开始</param>
    /// <returns></returns>
    public TKey select(int k)
    {
        return select(root, k).key;
    }

    private Node select(Node node, int k)
    {
        if (node == null)
            throw new IndexOutOfRangeException();
        int s = size(node.left);
        if (s == k)
            return node;
        else if (s < k)
            return select(node.right, k - s - 1);
        else
            return select(node.left, k);
    }

    public int size()
    {
        return size(root);
    }

    private int size(Node node)
    {
        return node?.size ?? 0;
    }

    /// <summary>
    /// [lo,hi]区间内的元素个数
    /// </summary>
    /// <param name="lo"></param>
    /// <param name="hi"></param>
    /// <returns></returns>
    public int size(TKey lo, TKey hi)
    {
        return size(root, lo, hi);
    }

    private int size(Node node, TKey lo, TKey hi)
    {
        if (node == null)
            return 0;
        int cnt = 0;
        int cmplo = compare(lo, node.key);
        int cmphi = compare(hi, node.key);
        if (cmplo <= 0 && cmphi >= 0)
            cnt += 1;
        if (cmplo < 0)
            cnt += size(node.left, lo, hi);
        if (cmphi > 0)
            cnt += size(node.right, lo, hi);
        return cnt;
    }
}