因為一句話,秒懂二叉樹旋轉

事情要從某天晚上買夜宵說起。買了香腸拿着吃,想着多年來一直沒搞懂的樹旋轉是不是應該看看,就點進某百科。

樹旋轉是在二叉樹中的一種子樹調整操作, 每一次旋轉並不影響對該二叉樹進行中序遍歷的結果。

中序遍歷!靈光一閃,好像很多東西都聯繫起來了!

為什麼是中序遍歷

中序遍歷是指按【左節點->父節點->右節點】的順序遍歷,這個內容能讓我們想起什麼?二叉排序樹。

二叉排序樹要求左節點的值小於父節點,右節點大於父節點。如果按照中序遍歷二叉排序樹,就能得到一個順序結果。

知道這一點有什麼用?在後續的旋轉過程中,我們可以根據二叉排序樹的父節點和子節點的大小關係來輔助理解旋轉。

舉個例子

現在我們有一顆二叉排序樹:

其中序遍歷的結果是 1、2、3、4、5、6、7、8、9。

以下先不考慮子樹的旋轉。

右旋

現在對根節點 5 執行右旋。右旋的時候,要選擇根節點 5 的左節點 3 作為新的根節點,稱為以節點 3 為轉軸。

為討論方便,把旋轉前 3 的右子樹稱為 A,5 的右子樹稱為 B。如下圖所示:

由於是右旋轉,當節點 3 處於根節點的時候,其左子樹的數必須仍然小於 3,又因為 A、5、B 都大於 3,所以 3 旋轉前的左子樹在旋轉後保持原樣,仍然是 3 的左子樹。

現在有兩個問題:

  1. A 放哪裡?
  2. 根節點 5 放在哪裡?

結合二叉排序樹來理解。由於 3 即將成為根節點,A 和原先根節點都大於 3,因此兩者都要放在旋轉後 3 的右子樹。那麼問題就轉化為:如何將 A、B、根節點 5 合併起來?

首先考慮根節點 5,先忽略 A。在自然旋轉後,節點 5 在 3 的右子樹,B 和 5 的關係不變,且 5 的左子樹必然為空。

現在考慮 A。由於旋轉前它就在 5 的左子樹裏面,所以必然小於 5。旋轉後 5 的左子樹為空,就可以直接把 A 作為 5 的左子樹。

旋轉完畢。用中序遍歷驗證,其結果是 1、2、3、4、5、6、7、8、9,結果不變。

右旋代碼實現

先不考慮子樹內旋轉的情況,便於理解。

package main

import "fmt"

type TreeNode struct {
	Left  *TreeNode
	Right *TreeNode
	Value int
}

// PutChild 用於簡化初始化
func (n *TreeNode) PutChild(child *TreeNode) {
	if child.Value < n.Value {
		n.Left = child
	} else if child.Value > n.Value {
		n.Right = child
	}
}

// RotateRight 右旋,參數先不使用轉軸
func RotateRight(root *TreeNode) {
	if root.Left == nil {
		return
	}

	// 把 A 備份出來
	a := root.Left.Right
	root.Left.Right = nil

	// 旋轉。3 原先在 5 左節點,現在讓 5 變成 3 的右節點
	newRoot := root.Left
	root.Left = nil
	newRoot.Right = root

	// 把 A 放回去
	root.Left = a
}

// PrintInorderIteration 中序遍歷迭代法
func PrintInorderIteration(root *TreeNode) {
	stack := make([]*TreeNode, 0)
	for len(stack) != 0 || root != nil {
		for root != nil {
			stack = append(stack, root)
			root = root.Left
		}

		root = stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		fmt.Println(root.Value)

		root = root.Right
	}
}

func main() {
	// 為了與圖對應,這裡多申請一個位置,但只從 1 開始初始化。
	nodes := make([]*TreeNode, 10)
	for i := 1; i < 10; i++ {
		nodes[i] = &TreeNode{Value: i}
	}

	nodes[5].PutChild(nodes[3])
	nodes[5].PutChild(nodes[8])

	nodes[3].PutChild(nodes[2])
	nodes[3].PutChild(nodes[4])

	nodes[2].PutChild(nodes[1])

	nodes[8].PutChild(nodes[7])
	nodes[8].PutChild(nodes[9])

	nodes[7].PutChild(nodes[6])

	fmt.Println("BEGIN")
	PrintInorderIteration(nodes[5])
	fmt.Println("END")

	RotateRight(nodes[5])

	fmt.Println("BEGIN")
	PrintInorderIteration(nodes[3])
	fmt.Println("END")
}

左旋

還是使用原始版本。這次根節點 5 的左子樹為 B,右節點 8 為轉軸,節點 8 的左子樹為 A(轉軸旋轉方向的內側)。

由於是左旋,轉軸 8 的左子樹 A 先不考慮。做自然旋轉。

由於子樹 A 之前是 8 的左子樹,說明子樹 A 上的節點都小於 8,因此旋轉後子樹 A 必然在新的根節點 8 的左側。

又由於旋轉後的節點 5 的右子樹必然為空,而子樹 A 在旋轉前就在 5 的右子樹,說明子樹 A 上的節點必然大於 5,因此旋轉後子樹 A 應作為 5 的右子樹。

旋轉完畢。用中序遍歷驗證,其結果是 1、2、3、4、5、6、7、8、9,結果不變。

代碼邏輯和右旋類似,不再重複。

子樹的旋轉

還是回到原先樹上。現在考慮以 3 為定點的子樹,執行右旋,此時轉軸為 2。

在旋轉過程中,與原先不同的地方在哪?

在於必須更新頂點的父節點。例如上圖,需要更新節點 5 的左節點為 2。按照前面右旋的代碼設計,當我們傳入的是節點 3 的時候,沒法更新節點 5。因此需要給節點引入一個新的字段 Parent,用於表示其父節點。

type TreeNode struct {
	Parent *TreeNode
	Left   *TreeNode
	Right  *TreeNode
	Value  int
}

在旋轉的時候,基於前面右旋的代碼,加上父節點的更新就可以了。有些步驟能合併,但為了容易理解,不予合併。

// PutChild 用於簡化初始化
func (n *TreeNode) PutChild(child *TreeNode) {
	if child.Value < n.Value {
		n.Left = child
	} else if child.Value > n.Value {
		n.Right = child
	} else {
		return
	}

	child.Parent = n
}

// RotateRight 右旋,參數先不使用轉軸
func RotateRight(root *TreeNode) {
	if root.Left == nil {
		return
	}

	// 把 A 備份出來,並取消雙向連線
	a := root.Left.Right
	root.Left.Right = nil
	if a != nil {
		a.Parent = nil
	}

	// 取到子樹根節點的父節點,並取消雙向連線
	rootParent := root.Parent
	root.Parent = nil
	// 如果 root 是整顆樹的根節點,無需調整
	if rootParent != nil {
		if rootParent.Value > root.Value {
			rootParent.Left = nil
		} else {
			rootParent.Right = nil
		}
	}

	// 旋轉。2 原先在 3 左節點,現在讓 3 變成 2 的右節點
	newRoot := root.Left
	root.Left = nil
	newRoot.Parent = nil

	newRoot.Right = root
	root.Parent = newRoot

	// 設置根節點的父節點
	if rootParent != nil {
		if rootParent.Value > newRoot.Value {
			rootParent.Left = newRoot
		} else {
			rootParent.Right = newRoot
		}
	}

	// 把 A 放回去
	root.Left = a
}

結尾

上文使用了二叉排序樹作為輔助理解的工具,順便一提,在想到二叉排序樹的時候,還應將其和二分搜索聯繫起來。

單看樹的旋轉,其實非常簡單。樹的旋轉實際上是為平衡二叉樹做鋪墊,因此下一篇將會把普通的二叉樹換成平衡二叉樹。