平衡二叉樹的 AVL 實現

上一篇【因爲一句話，秒懂二叉樹旋轉】把樹旋轉了解清楚，是爲這一篇平衡二叉樹準備的。

平衡二叉樹，就是在二叉樹的基礎上加上一個條件：對於任意節點，左子樹和右子樹的樹高之差不超過 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)
}

結尾

以前一直認爲二叉樹旋轉和平衡二叉樹很難，現在認真看的時候卻覺得其實也沒什麼。

接下去先用二叉樹的打印水一篇，然後就是紅黑樹了。

以前也一直覺得紅黑樹很難，但最近看了資料（特別是那本《算法》），覺得挺容易理解的，就乾脆一起寫出來吧。

平衡二叉樹的 AVL 實現

示例

可旋轉的二叉排序樹

添加 Height 參數

檢測樹平衡

旋轉時更新樹高

目前爲止的完整代碼

旋轉的問題

旋轉方式改進

結尾

【數據結構】Map （映射）的各種實現

在個人微信上與自搭服務端交互

工作的第二和第三年（201807~202005）

紅黑樹代碼實現及註釋

2-3 樹——學習紅黑樹的捷徑

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結