平衡二叉树的 AVL 实现

技术分享 2年前 (2022-06-02) 0 999+

关注

上一篇【因为一句话，秒懂二叉树旋转】把树旋转了解清楚，是为这一篇平衡二叉树准备的。

平衡二叉树，就是在二叉树的基础上加上一个条件：对于任意节点，左子树和右子树的树高之差不超过 1。

从实现的角度看，就是在已具备旋转功能的 Node 上增加一个 height 字段，并且在原先的代码上增加对 height 的操作。关键操作是判断左右子树的树高之差，根据树高之差选择需要执行的旋转。

示例

如图所示，这是一颗高为 3 的树。根节点左右两边的高度差为 1，满足平衡条件。

此时插入一个值为 1 的节点来破坏这个平衡。当节点 1 插入后，需要逐级往上更新节点的高度。

在高度更新后，发现根节点左右两边的高度差为 2，因此需要通过右旋调整平衡。节点 3 是转轴，按照旋转的规则执行。得到以下结果：

基本思路很简单。剩下的问题放到后面的内容处理。

以下内容会以可旋转的二叉排序树的代码（前篇文章的内容）为基础，添加 Height 这一属性。并根据树高的要求调整代码。

可旋转的二叉排序树

以下是原始代码，不难。与上一篇文章不同的是把 PutChild 改为 Insert，还有简化了旋转的代码。

package main  import "fmt"  type TreeNode struct { 	Parent *TreeNode  	Value  int  	Left   *TreeNode 	Right  *TreeNode }  // Insert 往树中合适位置插入节点 func Insert(root *TreeNode, value int) *TreeNode { 	if root == nil { 		return &TreeNode{Value: value} 	}  	if value < root.Value { 		root.Left = Insert(root.Left, value) 		root.Left.Parent = root 	} else if value > root.Value { 		root.Right = Insert(root.Right, value) 		root.Right.Parent = root 	} else { 		return root 	}  	return root }  // RotateRight 右旋 func RotateRight(root *TreeNode) *TreeNode { 	if root.Left == nil { 		return root 	}  	newRoot := root.Left 	tmp := newRoot.Right 	newRoot.Right = root 	root.Left = tmp  	if tmp != nil { 		tmp.Parent = root 	} 	newRoot.Parent = root.Parent 	root.Parent = newRoot  	return newRoot }  // RotateLeft 左旋 func RotateLeft(root *TreeNode) *TreeNode { 	if root.Right == nil { 		return root 	}  	newRoot := root.Right 	tmp := newRoot.Left 	newRoot.Left = root 	root.Right = tmp  	if tmp != nil { 		tmp.Parent = root 	} 	newRoot.Parent = root.Parent 	root.Parent = newRoot  	return newRoot }  // PrintTree 以树状形式打印树 func PrintTree(root *TreeNode) { 	// 这里先不管 }  func main() { 	var root *TreeNode 	root = Insert(root, 7) 	root = Insert(root, 3) 	root = Insert(root, 2) 	root = Insert(root, 5) 	root = Insert(root, 8) 	PrintTree(root) 	fmt.Println("------------")  	root = RotateLeft(root) 	PrintTree(root) 	fmt.Println("------------")  	root = RotateRight(root) 	PrintTree(root) 	fmt.Println("------------") }

添加 Height 参数

type TreeNode struct { 	Parent *TreeNode 	 	Value  int 	Height int  	Left   *TreeNode 	Right  *TreeNode }  func NewTreeNode(value int) *TreeNode { 	return &TreeNode{Value: value, Height: 1} }

因为每次插入的节点都是作为叶子节点，所以新节点的树高都为 1。这里新增 New 函数，在初始化时自动指定。

检测树平衡

在修改的时候才需要检测树平衡，因此需要关注三个操作：插入、旋转、删除。这里忽略删除的部分，仅关注插入和旋转。

原始代码：

func Insert(root *TreeNode, value int) *TreeNode { 	if root == nil { 		return &TreeNode{Value: value} 	}  	if value < root.Value { 		root.Left = Insert(root.Left, value) 		root.Left.Parent = root 	} else if value > root.Value { 		root.Right = Insert(root.Right, value) 		root.Right.Parent = root 	} else { 		return root 	}  	return root }

每次插入一个叶子节点，有可能使其父节点树高增加，因此必须更新其父节点的树高。接着由于父节点树高的更新，需要再更新上一层树高，直到根节点。

由于 Insert 是递归调用，在递归下面写插入完成后的代码。分为两部分：1. 更新当前节点的树高；2. 如果左右子树树高相差超过 1，则通过旋转平衡该树。

首先第一部分，更新树高。

由于左右子树高度不一定一致，所以要取较高的那一颗子树的树高，加上 1 就是以当前节点作为根节点的子树的高度。

func max(a, b int) int { 	if a > b { 		return a 	} 	return b }  // GetHeight 用来处理节点为 nil 的情况 func GetHeight(node *TreeNode) int { 	if node == nil { 		return 0 	}  	return node.Height }  func Insert(root *TreeNode, value int) *TreeNode { 	// ...  	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1  	return root }

接着第二步，判断树平衡。

引入一个函数来获取平衡因子。当平衡因子小于 -1 的时候，表示右边的子树比左边高大于 1，此时应该左旋。反之，平衡因子大于 1 的时候表示应该右旋。

// GetBalanceFactor 获取平衡因子 func GetBalanceFactor(node *TreeNode) int { 	if node == nil { 		return 0 	}  	return GetHeight(node.Left) - GetHeight(node.Right) }  func Insert(root *TreeNode, value int) *TreeNode { 	// ...  	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1  	bf := GetBalanceFactor(root) 	if bf < -1 { // 应该左旋 		root = RotateLeft(root) 	} else if bf > 1 { // 应该右旋 		root = RotateRight(root) 	} else { 		// do nothing 	}  	return root }

旋转时更新树高

这里要先更新原先 root 节点的树高，因为旋转后它是 newRoot 的子节点。总是要按照先子节点再父节点的顺序更新树高。

另外由于 tmp 子树本身没有修改，因此不需要更新树高。

// RotateRight 右旋 func RotateRight(root *TreeNode) *TreeNode { 	// ...  	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1 	newRoot.Height = max(GetHeight(newRoot.Left), GetHeight(newRoot.Right)) + 1  	return newRoot }  // RotateLeft 左旋 func RotateLeft(root *TreeNode) *TreeNode { 	// ...  	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1 	newRoot.Height = max(GetHeight(newRoot.Left), GetHeight(newRoot.Right)) + 1  	return newRoot }

目前为止的完整代码

在进入下一个阶段前，有必要浏览一遍当前的完整代码。

package main  import "fmt"  type TreeNode struct { 	Parent *TreeNode  	Value  int 	Height int  	Left  *TreeNode 	Right *TreeNode }  func NewTreeNode(value int) *TreeNode { 	return &TreeNode{Value: value, Height: 1} }  // Insert 往树中合适位置插入节点 func Insert(root *TreeNode, value int) *TreeNode { 	if root == nil { 		return &TreeNode{Value: value} 	}  	if value < root.Value { 		root.Left = Insert(root.Left, value) 		root.Left.Parent = root 	} else if value > root.Value { 		root.Right = Insert(root.Right, value) 		root.Right.Parent = root 	} else { 		return root 	}  	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1  	bf := GetBalanceFactor(root) 	if bf < -1 { // 应该左旋 		root = RotateLeft(root) 	} else if bf > 1 { // 应该右旋 		root = RotateRight(root) 	} else { 		// do nothing 	}  	return root }  func max(a, b int) int { 	if a > b { 		return a 	} 	return b }  // GetHeight 用来处理节点为 nil 的情况 func GetHeight(node *TreeNode) int { 	if node == nil { 		return 0 	}  	return node.Height }  // GetBalanceFactor 获取平衡因子 func GetBalanceFactor(node *TreeNode) int { 	if node == nil { 		return 0 	}  	return GetHeight(node.Left) - GetHeight(node.Right) }  // RotateRight 右旋 func RotateRight(root *TreeNode) *TreeNode { 	if root.Left == nil { 		return root 	}  	// 旋转 	newRoot := root.Left 	tmp := newRoot.Right 	newRoot.Right = root 	root.Left = tmp  	// 更新节点的父节点信息 	if tmp != nil { 		tmp.Parent = root 	} 	newRoot.Parent = root.Parent 	root.Parent = newRoot  	// 更新树高 	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1 	newRoot.Height = max(GetHeight(newRoot.Left), GetHeight(newRoot.Right)) + 1  	return newRoot }  // RotateLeft 左旋 func RotateLeft(root *TreeNode) *TreeNode { 	if root.Right == nil { 		return root 	}  	// 旋转 	newRoot := root.Right 	tmp := newRoot.Left 	newRoot.Left = root 	root.Right = tmp  	// 更新节点的父节点信息 	if tmp != nil { 		tmp.Parent = root 	} 	newRoot.Parent = root.Parent 	root.Parent = newRoot  	// 更新树高 	root.Height = max(GetHeight(root.Left), GetHeight(root.Right)) + 1 	newRoot.Height = max(GetHeight(newRoot.Left), GetHeight(newRoot.Right)) + 1  	return newRoot }  func PrintTree(root *TreeNode) { }  func main() { 	var root *TreeNode 	root = Insert(root, 7) 	root = Insert(root, 3) 	root = Insert(root, 2) 	root = Insert(root, 5) 	root = Insert(root, 8) 	PrintTree(root) 	fmt.Println("------------")  	root = Insert(root, 6) 	PrintTree(root) 	fmt.Println("------------") }

旋转的问题

与最开始示例不同的是，上面的代码最后插入的是 6。执行这些代码，发现得到的仍然是一颗不平衡的树。

为什么？

用图来解释：

这是原始图，插入节点 6 后得到：

此时对于节点 7，左子树比右子树高 2。因此需要右旋，根据规则右旋后得到：

为什么会这样？

从第二张图可以看到，之所以平衡被打破，是因为 A 的高度发生变化，导致节点 3 的高度变化。当右旋开始时，这个打破平衡的 A 被抽离了。

抽离后节点 3 的高度变回 2，也就是说根节点左子树的高度为 2。如果执行右旋，那么根节点的左子树的高度必定会减 1，变成 1。

不管原先根节点右子树的树高是多少，在旋转后树高为 2 的部分必然要挂到原先根节点 7 上。此时根节点的右子树高度必大于 2。

上图把节点 8 隐藏起来，可以直观地看到问题。

也就是说，继续按照之前的方式，旋转后必处于不平衡状态。而且从这个状态出发，也只能执行左旋，旋转后回到原来的样子。进入死循环。

怎么解决？

根据上面的描述，右旋是必然要做的，并且右旋时必然使左子树高度减 1。

要解决这个问题，需要让根节点的左子树去掉其右子树剩下的部分的高度增加，然后再右旋。

有没有办法？

有，让左子树左旋就行。

根据规则旋转后：

接着，对根节点执行右旋：

去掉 A 的部分后，左子树的高度仍然为 3，右旋后为 2。

平衡了。

综上，当一个节点的左右子树不平衡时，要分两步判断：

左子树高还是右子树高
不平衡是由子树的左子树还是右子树引起的

如果是左子树的右子树引起的，则子树需要先旋转。同理，如果是右子树的左子树引起的，子树也要先旋转。（发现没？两者都是在旋转曲线的内侧）

旋转方式改进

从前面的内容可以总结出旋转有四种类型：

左子树的左子树引起的失衡，用 LL（Left-Left）表示；
左子树的右子树引起的失衡，用 LR（Left-Right）表示；
右子树的左子树引起的失衡，用 RL（Right-Left）表示；
右子树的右子树引起的失衡，用 RR（Right-Right）表示。

在 Insert 的时候，要区分这四种情况。

func Insert(root *TreeNode, value int) *TreeNode { 	// ...  	bf := GetBalanceFactor(root) 	if bf < -1 { // 应该左旋 		if value < root.Right.Value { // 在右子树的左子树上 			root = RLRotation(root) 		} else { // 在右子树的右子树上 			root = RRRotation(root) 		} 	} else if bf > 1 { // 应该右旋 		if value < root.Left.Value { // 在左子树的左子树上 			root = LLRotation(root) 		} else { // 在左子树的右子树上 			root = LRRotation(root) 		} 	} else { 		// do nothing 	}  	return root }  func LLRotation(root *TreeNode) *TreeNode { 	return RotateRight(root) }  func LRRotation(root *TreeNode) *TreeNode { 	root.Left = RotateLeft(root.Left) 	return RotateRight(root) }  func RRRotation(root *TreeNode) *TreeNode { 	return RotateLeft(root) }  func RLRotation(root *TreeNode) *TreeNode { 	root.Right = RotateRight(root.Right) 	return RotateLeft(root) }