技術文章 » 平衡二元樹的 AVL 實現

平衡二元樹的 AVL 實現

2022-06-01 21:00:35

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

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

平衡了。

綜上，當一個節點的左右子樹不平衡時，要分兩步判斷：

  1. 左子樹高還是右子樹高
  2. 不平衡是由子樹的左子樹還是右子樹引起的

如果是左子樹的右子樹引起的，則子樹需要先旋轉。同理，如果是右子樹的左子樹引起的，子樹也要先旋轉。（發現沒？兩者都是在旋轉曲線的內側）

旋轉方式改進

從前面的內容可以總結出旋轉有四種型別：

  1. 左子樹的左子樹引起的失衡，用 LL（Left-Left） 表示；
  2. 左子樹的右子樹引起的失衡，用 LR（Left-Right） 表示；
  3. 右子樹的左子樹引起的失衡，用 RL（Right-Left） 表示；
  4. 右子樹的右子樹引起的失衡，用 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)
}

結尾

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

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

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