数据结构系列：平衡二叉树之 AVL 树

1 前言

当我们查询特定元素时，二叉搜索树在大多情况下的性能会比线性表强的多。但并不总是这样，极端情况下，二叉搜索树也会退化成线性表，比如我们插入从小到大的序列（从大到小的序列），这样会导致二叉搜索树中的每个节点的左孩子（右孩子）都为空。可不可以避免上述这种情况呢？可以。我们只需要在添加元素后调整树的结构，让每一个节点的左右孩子尽量都不为空即可，平衡二叉树中的平衡就是这个意思。下面我们介绍一种平衡二叉树：AVL 树。

2 AVL 树简介

AVL 树得名于它的发明者 Adelson-Velsky 和 Landis， 它是平衡二叉树的一种（后续将要介绍的红黑树也是平衡二叉树）。那什么是平衡二叉树？平衡又是什么意思？

平衡二叉树没有绝对的定义，在一棵二叉树中，所有节点左子树和右子树的高度越接近，它就越平衡。也就是说，平衡与左右子树的高度差有关系。为了计算高度差，平衡二叉树的节点比二叉搜索树的节点会多一个高度属性。

// 内部类
private class AVLNode<E> extends Node<E>{
	......
	public int height = 1;	// 默认为 1 的原因是：新添加的节点一定是叶子节点，所以高度为 1。
	......
}

在AVL 树中，平衡是这样定义的：任一节点的左右子树的高度差的绝对值不能超过 1。也就是说，如果在添加元素或者删除元素的过程中，某一个节点违反了这个条件，那便需要经过调整恢复到平衡状态。

3 AVL 树如何保证平衡（代码实现）

AVL 树也是一棵二叉搜索树，因此，我们可以直接继承它，复用它的代码。与之前的代码有一点不同，AVL 树会在添加删除元素后做进一步处理，添加 afterAdd() 方法 和 afterRemove() 方法。在这两个两个方法中，我们需要知道是哪一个节点导致二叉搜索树失衡，怎么找到这个节点，以及如何恢复平衡。对于这三个问题，添加元素和删除元素有着各自不同的逻辑，下面我们一一介绍。

3.1 添加元素失衡后如何恢复平衡

在二叉搜索树中每次添加完元素以后，我们需要对节点做一次扫描，目的是找出不平衡的节点并让其恢复平衡。

由于我们添加的节点最终会成为叶子节点，因此，高度可能变化的节点只能是该节点的父节点和祖先节点（从父节点的父节点开始，一直往上的节点都是祖先节点），但高度变化不一定就会失衡。其中，父节点就一定不会失衡，这点可以仔细思考一下（但父节点的高度可能会变）。只有祖先节点可能会失衡。最差的情况是由于一个祖先节点失衡，可能导致所有往上的祖先节点都失衡！

怎么找到失衡的祖先节点呢？非常简单，沿着父节点一直往上走，哪个祖先节点的左右子树的高度差的绝对值超过 1 ，谁就失衡。此时，我们需要调整以该失衡节点为根的二叉树，让这棵子树回到平衡状态（调整以该失衡节点为根的二叉搜索树，使其高度减 1）。同时，我们也需要更新这条路径上各个节点的高度值。

前面我们提到最差的情况会因为一棵子树高度变化而导致所有祖先节点失衡，那是不是应该对所有失衡的祖先节点进行调整呢？答案是否定的。我们添加一个节点，失衡的祖先节点的高度最多加 1。我们只需要对离添加节点距离最近、且失衡的那一个祖先节点为根的二叉搜索树进行调整，让其恢复平衡状态即可。这样，再上面的祖先节点的高度属性也会恢复到未添加该节点之前的状态。

怎么调整失衡节点的状态呢？这里一共包含四种情况（以下资料参考自小码哥的视频教程）。

下面四张 PPT 中，一共包含了 LL, RR, LR, RL 的四种情况，图中从左到右也画出了对应情况的处理方法。其中字母 L 代表左，R 代表右。比如，LL 代表在节点 g (grandparent) 的左孩子的左子树上添加一个节点导致了失衡，并且节点 g 是第一个失衡的祖先节点。 RR, LR, RL 也是类似。（图中的字母 p 为 parent 的简写，n 为 node 的简写。T0、T1、T2、T3 是二叉树，其中 T0 和 T1 的高度相等，T2 和 T3 的高度相等。红色的节点为添加的节点，黑红黑三根虚线为各棵子树高度的参考线）

（1）LL：右旋转

（2）RR：左旋转

（3）LR：RR - 左旋转，LL - 右旋转

（4）RL：LL - 右旋转，RR - 左旋转

我们需要熟知右旋转和左旋转，这没有什么技巧可言，花时间看懂熟悉即可。 旋转过后，我们需要更新一些节点的 parent 属性和高度。

下面我将进行代码实现。

我们首先需要创建一个 afterAdd 方法，这个方法将新添加的叶子节点传进去。我们会根据这个叶子节点一直往上（父节点）寻找，判断节点是否失衡。如果没有失衡，我们需要更新节点的高度值；如果失衡了，我们会判断此时的失衡状态是上述四种（LL, RR, LR, RL）的哪一种状态，然后对应进行调整。下面，我们给出整个方法的大体框架。

	protected void afterAdd(Node<E> node){
		
		Node<E> parentNode = node.parent;	// 这里我们可以进行优化，如果 parentNode 的高度值并未改变，那么就不用继续往上找了。
		
		while(parentNode != null){	// 一直往上，直到根节点为止
		
			if(isBalanced(parentNode)){
				updateHeight(parentNode);	// 节点没有失衡，我们需要检查节点的高度值; 	
			}else{
				rebalance(parentNode);
				break;
			}
			parentNode = parentNode.parent;
		}
		
	}

上述代码中的三个方法，isBalanced 、updateHeight 和 rebalance 对应我们要解决的 3 个问题：（1）如何判断一棵二叉树是否失衡；（2）如何更新节点的高度；（3）如何让一棵二叉树恢复至平衡状态。我们一一来实现这些方法。

3.1.1 如何实现 isBalanced 方法

根据 AVL 树的定义，如果一个节点的左右子树高度差的绝对值超过 1，则被认为是失衡的。因此，我们只需要知道左右子树的高度，便可以完成这个方法。

	private boolean isBalanced(Node<E> root){
		int leftHeight = root.left == null ? 0: root.left.height;
		int rightHeight = root.right == null ? 0: root.right.height;
		return Math.abs(leftHeight - rightHeight) <= 1;
	}

3.1.2 如何实现 updateHeight 方法

实现这个方法的逻辑和 isBalanced 方法的逻辑类似，都需要获取左右子树的高度。

	private void updateHeight(Node<E> root){
		int leftHeight = root.left == null ? 0: root.left.height;
		int rightHeight = root.right == null ? 0: root.right.height;
		root.height = Math.max(leftHeight, rightHeight) + 1;
	}

到这里有同学会发现上述的两个方法有重复代码，因此，我们可以将代码重构一下。

我们直接在 AVLNode 类中封装 leftHeight 和 rightHeight 方法，用于获取当前节点左右子树的高度值。进而我们可以在里面封装**计算平衡因子（左右子树的高度差）**的方法用于 isBalanced 方法中。同时，获取了左右子树的高度也可以进一步更新当前节点高度。

class AVLNode<E> extends Node<E>{
	......
	// 获取左子树的高度	
	public int leftHeight(){
		int leftHeight = this.left == null ? 0: this.left.height;
		return leftHeight;
	}
	
	// 获取右子树的高度	
	public int rightHeight(){
		int rightHeight = this.right == null ? 0: this.right.height;
		return rightHeight;
	}
	
	// 计算节点的平衡因子
	public int balanceFactor(){
		return leftHeight() - rightHeight();
	}
	
	// 更新节点的高度值
	public void updateHeight(){
		this.height = Math.max(leftHeight, rightHeight) + 1;
	}
	......
}
	

因此，上述的 isBalanced 方法和 updateHeight 方法可以简化成：

	private boolean isBalanced(Node<E> node){
		return Math.abs(((AVLNode)node).balanceFactor()) <= 1;
	}

	private void updateHeight(Node<E> node){
		((AVLNode)node).updateHeight();
	}

3.1.3 如何实现 rebalance 方法

这部分的实现最需要我们进行深思，具体的实现过程需要参照上面四张 PPT。

首先我们需要解决的问题是：当前的失衡状态属于四张 PPT 中的哪一种？对于这个问题，我们需要明确 g、p、n 三个节点的关系，如果 p 是 g 的左孩子，n 又是 p 的左孩子，这种情况属于 LL 情况，其他情况也是同理。

但是，怎么确定 p 和 n 是父节点的左孩子还是右孩子呢？理清这个问题，我们需要明白失衡的原因：由于添加了一个新的叶子节点，所以导致以 g 为根的二叉树失衡了。要恢复平衡，我们需要调整 g 的高度较高的子树（左子树或者右子树），降低它的高度，这样，以 g 为根节点的二叉树的高度也就降下来了。

因此，p 和 n 是左孩子还是右孩子，完全取决于它在哪一棵高度更高的子树上。因此，我们需要一个方法获取一个节点更高的那棵子树。

class AVLNode<E> extends Node<E>{
	......
	// 获取高度更高的子树
	private Node<E> tallerChild(){
	
		int leftH = leftHeight();
		int rightH = rightHeight();
		if(leftH > rightH){
			
			return left;
		}else if(leftH < rightH){
			
			return right;
		}else{
			
			return isLeftChild() ? left : right;	// 如果当前节点左右子树的高度相等，我们则根据当前节点是左孩子还是右孩子来判断。如果是左孩子，则返回左孩子，反之，返回右孩子。这样做的原因是为了转换成 LL 或者 RR 情况，恢复平衡时只需要旋转一次。
		}
	}
	......
}

自此，我们可以大致确认 rebalance 方法的逻辑了。

public void rebalance(Node<E> grand){
	Node<E> parent = ((AVLNode)grand).tallerChild();
	Node<E> node = ((AVLNode)parent).tallerChild();
	
	// 分情况
	if(parent.isLeftChild()){	// L-
		
		if(node.isLeftChild()){	// LL
		
			rotateRight(grand);		// 对 grand 进行右旋转
		}else{				    // LR
		
			rotateLeft(parent);		// 先对 parent 进行左旋转变成 LL 情况
			rotateRight(grand);		// 再对 grand 进行右旋转
		}
	}else{						// R-				
	
		if(node.isLeftChild()){ // RL
		
			rotateRight(parent);	// 先对 parent 进行右旋转变成 RR 情况
			rotateLeft(grand);		// 再对 grand 进行左旋转
		}else{					// RR
		
			rotateLeft(grand);			// 对 grand 进行左旋转
		}
	}
	
}

重点来了，如何实现 rotateLeft 方法和 rotateRight 方法。实际上两者的实现逻辑是一样的，我们以 ratateLeft 方法为例。

参照下图（这张图上面也有，为了让大家更方便看到放这儿是一个更好的选择），我们首先要（1）更改节点 g 和 p 的指针；（2）然后更新 T1、p 和 g 的 parent 属性；（3）最后我们再更新 g 和 p 的高度。这个顺序能调换吗？（3）必须在（1）的后面，（2）更新 parent 属性可以随意。

private void rotateLeft(Node<E> grand){
	
	// 根据节点 grand 获取节点 parent
	Node<E> parent = grand.right;
	
	// 调整指针指向
	grand.right = parent.left;
	parent.left = grand;

	// 更改节点的 parent 属性
	// 1、维护 parent 的父节点 
	grand.parent = parent.parent;
	if(grand.isLeftChild()){
		grand.parent.left = parent;
	}else if(grand.isRightChild()){
		grand.parent.right = parent;
	}else{
		root = parent;
	}

	// 2、维护 grand 右孩子的父节点
	Node<E> child = grand.right;
	if(child != null){
		child.parent = grand;
	}
	
	// 3、维护 grand 的父节点，这个要放在 1 的后面，维护 parent 的父节点需要依赖 grand 原始的父节点
	grand.parent = parent;
	
	// 更新节点的高度，两者的顺序不能反
	updateHeight(grand);
	updateHeight(parent);
}

同理，我们也可以建立 rotateRight 的逻辑：

private void rotateRight(Node<E> grand){
	
	// 根据节点 grand 获取节点 parent
	Node<E> parent = grand.left;
	
	// 调整指针指向
	grand.left = parent.right;
	parent.right = grand;

	// 更改节点的 parent 属性
	// 1、维护 parent 的父节点 
	g.parent = p.parent;
	if(grand.isLeftChild()){
		grand.parent.left = parent;
	}else if(grand.isRightChild()){
		grand.parent.right = parent;
	}else{
		root = parent;
	}

	// 2、维护 parent 右孩子的父节点
	Node<E> child = grand.left;
	if(child != null){
		child.parent = grand;
	}
	
	// 3、维护 grand 的父节点，这个要放在 1 的后面，维护 parent 的父节点需要依赖 grand 原始的父节点
	grand.parent = parent;
	
	// 更新节点的高度，两者的顺序不能反
	updateHeight(grand);
	updateHeight(parent);
}

仔细观察我们可以发现，两个方法中从更新节点的 parent 属性开始，代码有大量重复的地方，除了 child 变量的指向不同，其余都相同，因此，我们可以把这块代码抽出来形成一个私有的方法 afterRotate：

private void afterRotate(Node<E> grand, Node<E> parent, Node<E> child){

	// 更改节点的 parent 属性
	// 1、维护 parent 的父节点 
	grand.parent = parent.parent;
	if(grand.isLeftChild()){
		grand.parent.left = parent;
	}else if(grand.isRightChild()){
		grand.parent.right = parent;
	}else{
		root = parent;
	}

	// 2、维护 parent 右孩子的父节点
	if(child != null){
		child.parent = grand;
	}
	
	// 3、维护 grand 的父节点，这个要放在 1 的后面，维护 parent 的父节点需要依赖 grand 原始的父节点
	grand.parent = parent;
	
	// 更新节点的高度，两者的顺序不能反
	updateHeight(grand);
	updateHeight(parent);	
}

这样 rotateLeft 方法和 rotateRight 方法可以简化成如下形式：

private void rotateLeft(Node<E> grand){
	
	// 根据节点 grand 获取节点 parent，进而获取 child
	Node<E> parent = grand.right;
	Node<E> child = parent.left;
	// 调整指针指向
	grand.right = parent.left;
	parent.left = grand;
	
	afterRotate(grand, parent, child);
}

private void rotateRight(Node<E> grand){
	
	// 根据节点 grand 获取节点 parent，进而获取 child
	Node<E> parent = grand.left;
	Node<E> child = parent.right;
	
	// 调整指针指向
	grand.left = parent.right;
	parent.right = grand;
	
	afterRotate(grand, parent, child);
}

添加元素失衡后的恢复平衡过程已经完成，总体上看其实并不复杂，只是过程长一些。下面我们看看删除元素失衡后如何恢复平衡。

3.2 删除元素失衡后如何恢复平衡

删除节点和添加节点类似，也需要在删除动作完成后添加 afterRemove 方法。我们需要检测哪些节点失衡了，将其恢复至平衡状态即可。问题来了，我们应该从哪儿开始寻找失衡节点呢？找到失衡的节点以后，是不是也像 afterAdd 方法那般仅仅需要调整一个节点的平衡状态就可以了呢？

3.2.1 从哪儿开始寻找失衡节点

我们先回答第一个问题：从哪儿开始寻找失衡节点？对于删除操作，**无论删除的是叶子节点还是非叶子节点，最终删除的位置一定是某一个叶子节点。**这点可以回顾删除元素的 3 种情况，如果删除的是非叶子节点，最终也是用某一个叶子节点替代，与该叶子节点相关的子树高度也可能会减少。因此，我们需要从被删除的叶子节点开始，一直往上寻找失衡节点。

3.2.2 如何 rebalance

第二个问题：删除导致的失衡和添加一样，仅仅只需要 rebalance 一次吗？（下图中的 break 语句表明添加导致的失衡只需要 rebalance 一次）对于这个问题，我们需要明白为什么添加导致的失衡只需要 rebalance 一次。

添加的元素如果导致失衡，该元素所在子树的高度会加 1，我们往上遇到第一个不平衡的节点时，将其所在子树恢复至平衡以后高度又会减 1，至此，整棵树再往上的节点的高度状态和未添加该节点之前是一样的。

现在我们分析删除元素的情况。当我们删除一个元素，如果导致了某棵子树失衡（失衡说明左右子树高度差的绝对值小于 1，该子树的高度可能不变，也可能减 1），对其恢复平衡以后，整棵子树的高度必定会减 1。这种变化可能会导致往上的父节点继续失衡。因此，删除节点导致的失衡需要一直往上 rebalance，直到根节点。

下面，我们可以得出方法 afterRemove 的逻辑了：

protected void afterRemove(Node<E> node){
	// 一直调整，直至根节点
	while(node != null){

        if(isBalanced(node)){
            // 如果平衡，则更新祖先节点的高度即可
            updateHeight(node);

        }else{
            // 恢复平衡，则根据情况旋转
            rebalance(node);    //  node 为不平衡的父节点
        }

        node = node.parent;
    }
}

我们需要注意的是在之前删除元素的过程中，我们并未改变被删除节点的 parent 属性，这才可以直接将被删除节点直接传入上述的 afterRemove 方法中。

4 总结

AVL 树通过在添加元素后做进一步处理，使得它的查询性能比二叉搜索树要好。而查询的性能也会直接影响添加元素和删除元素的性能，因此，整体上AVL 树比二叉搜索树要更好。可当我们在 AVL 树中删除一个节点时，我们需要进行 O(logn) 次的旋转调整，这在一定程度上影响了删除元素的效率。而接下来我们介绍的红黑树，它在这一点上进行了改进。相较于 AVL 树，红黑树是一种统计性能更优的二叉树，详情请见数据结构系列：红黑树。