源码

来自半年后的说明：我把代码转放到 gists 里了，当时代码风格受 Python 影响较大，这里就不再更改了，毕竟大家本地都有自己的格式化工具。

点击这里前往Github获取本文源码。

介绍

AVL树（Adelson-Velsky and Landis Tree）是最早被发明的自平衡二叉查找树，它能保证查找、插入和删除在平均和最坏情况下的时间复杂度都是O(log n)。其内部的原理就是增加和删除的时候都会借助一次或多次旋转操作来实现树的重新平衡。下面是几个概念：

  • Height，高度，是当前节点一共有几层子节点，所以单个叶子节点的高度是0。
  • Balance Factor，平衡因子，是当前节点的左子树高度减去右子树高度。

当平衡因子处于[-1, 1]区间时，我们认为这棵树是平衡的，否则就是不平衡状态，需要通过一次或多次旋转使其重新平衡。

如果你还不知道什么是二叉查找树，请看点这里看我写的上一篇文章。

实现

逐个函数分析。

创建节点

同BST中的代码基本相同，但是它多了一个height属性，用来储存当前节点的高度。代码如下：

1
2
3
function createNode(val) {
    return { val, left: undefined, right: undefined, height: 0 }
}

高度

由于节点可能是undefined或null，所以不可以直接用node.height的形式获取高度，这里写一个函数height来帮助我们进行判空。

同时，我们再写一个函数calcHeight，通过子节点的高度计算当前节点的高度。代码如下：

1
2
3
4
5
6
7
function height(node) {
    return node ? node.height : -1
}

function calcHeight(node) {
    return Math.max(height(node.left), height(node.right)) + 1
}

平衡因子

左子树高度减去右子树高度，当node为空的时候，返回0，代码如下：

1
2
3
4
5
6
function factor(node) {
    if (!node) {
        return 0
    }
    return height(node.left) - height(node.right)
}

查找

查找因为不涉及更改节点，所以与普通的BST相同，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
function contains(node, val) {
    if (!node) {
        return false
    }

    if (val < node.val) {
        return contains(node.left, val)
    } else if (val > node.val) {
        return contains(node.right, val)
    }
    
    return true
}

旋转

旋转就是AVL树的核心操作了，一共有四种旋转模式，对应着四种不同的不平衡情况。

左单旋转

当node.left.left被进行了一次插入操作，导致这棵树不平衡时，需要进行左单旋转，过程如下：

左单旋转

分析：

  1. 由于插入了节点x，使得原本以k1为根节点的AVL树不再平衡。
  2. 考虑将A树上升一层，将C树下移一层，通过让k2取代k1成为新的根节点可以做到。
  3. 那么B树放到哪里？根据二叉搜索树的定义，我们知道，对于任意B树中的节点m，都有m > k2 && m < k1，所以它应该被放置在k2之右、k1之左，所以就放到了图示的位置。

此外，还需考虑高度的重新计算问题，这里先贴上代码：

1
2
3
4
5
6
7
8
9
10
function rotateWithLeftChild(k1) {
    const k2 = k1.left
    k1.left = k2.right
    k2.right = k1

    k1.height = calcHeight(k1)
    k2.height = calcHeight(k2)

    return k2
}

通过图我们可以知道，k1的左右子树高度（B 和 C）是不变的，所以先算k1，然后计算k2，就可以算出正确的高度了（至于子节点高度，每次插入都会重置）。

右单旋转

当node.right.right被进行了一次插入操作，导致这棵树不平衡时，需要进行右单旋转，过程如下：

右单旋转

基本和左单旋转相同，这里不多做解释，直接贴上代码：

1
2
3
4
5
6
7
8
9
10
function rotateWithRightChild(k1) {
    const k2 = k1.right
    k1.right = k2.left
    k2.left = k1

    k1.height = calcHeight(k1)
    k2.height = calcHeight(k2)

    return k2
}

左双旋转

当node.left.right被进行了一次插入操作，导致这棵树不平衡时，需要进行左双旋转，过程如下

左双旋转

分析：

  1. 由于插入了节点x，使得原以k3为根节点的AVL树不再平衡。
  2. 我们先将k1进行右单旋转，得到中间的图，从而使得整个树有一种斜向上的形状。
  3. 所以再将k3进行左单旋转，得到最后的图，实现重新平衡。

因为我们将其拆解成了两步单旋转，所以代码反而更简单，如下：

1
2
3
4
function doubleWithLeftChild(k3) {
    k3.left = rotateWithRightChild(k3.left)
    return rotateWithLeftChild(k3)
}

右双旋转

当node.right.left被进行了一次插入操作，导致这棵树不平衡时，需要进行右双旋转，过程如下：

右双旋转

代码实现如下：

1
2
3
4
function doubleWithRightChild(k3) {
    k3.right = rotateWithLeftChild(k3.right)
    return rotateWithRightChild(k3)
}

平衡

我们要根据情况的不同选择不同的旋转函数，所以这里单独用一个函数balance用来平衡这个树，代码如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
function balance(node) {
    if (!node) {
        return node
    }
    if (factor(node) > 1) {
        // Left
        if (factor(node.left) < 0) {
            node = doubleWithLeftChild(node)
        } else {
            node = rotateWithLeftChild(node)
        }
    } else if (factor(node) < -1) {
        // Right
        if (factor(node.right) > 0) {
            node = doubleWithRightChild(node)
        } else {
            node = rotateWithRightChild(node)
        }
    }
    // Anyway, reset the height.
    node.height = calcHeight(node)
    return node
}

分析：

  1. 如果给的节点是空，返回它，不进行任何操作。
  2. 如果平衡因子大于1，也就是node.left比node.right高出1层多，执行以下逻辑：
    1. 如果node.left的平衡因子小于0，也就是node.left.right更高，执行双旋转。
    2. 否则，就是node.left.left更高，执行单旋转。这里不可能两个子树一样高，因为刚打破平衡时这棵树就要被重新调整了。
  3. 如果平衡因子小于-1，也就是node.right比node.left高出1层多，执行以下逻辑：
    1. 如果node.right的平衡因子大于0，也就是node.right.left更高，执行双旋转。
    2. 否则，就是node.right.right更高，执行单旋转。这里也是不可能一样高，原因同上。
  4. 最后重新设置以下node.height，返回node，结束操作。由于插入和删除操作都依赖于本函数，所以每次操作都可以正确地更新每个节点的高度。

插入

除了最后一句把AVL树重新平衡外，其它与普通的BST相同，不多做解释：

1
2
3
4
5
6
7
8
9
10
11
function insert(node, val) {
    if (!node) {
        return createNode(val)
    }
    if (val < node.val) {
        node.left = insert(node.left, val)
    } else if (val > node.val) {
        node.right = insert(node.right, val)
    }
    return balance(node)
}

删除

同样，除了最后一句重新平衡了这个树，其它与普通BST相同，不多解释：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
function remove(node, val) {
    if (!node) {
        return node
    }
    if (val < node.val) {
        node.left = remove(node.left, val)
    } else if (val > node.val) {
        node.right = remove(node.right, val)
    } else if (node.left && node.right) {
        // Find minimum value
        let min = node.right
        while (min.left) {
            min = min.left
        }
        node.val = min.val
        node.right = remove(node.right, node.val)
    } else {
        node = node.left || node.right
    }
    return balance(node)
}

参考

数据结构与算法分析