因为一句话，秒懂二叉树旋转

2022 年 5 月 18 日
筆記
算法与数据结构

事情要从某天晚上买夜宵说起。买了香肠拿着吃，想着多年来一直没搞懂的树旋转是不是应该看看，就点进某百科。

树旋转是在二叉树中的一种子树调整操作，每一次旋转并不影响对该二叉树进行中序遍历的结果。

中序遍历！灵光一闪，好像很多东西都联系起来了！

为什么是中序遍历

中序遍历是指按【左节点->父节点->右节点】的顺序遍历，这个内容能让我们想起什么？二叉排序树。

二叉排序树要求左节点的值小于父节点，右节点大于父节点。如果按照中序遍历二叉排序树，就能得到一个顺序结果。

知道这一点有什么用？在后续的旋转过程中，我们可以根据二叉排序树的父节点和子节点的大小关系来辅助理解旋转。

举个例子

现在我们有一颗二叉排序树：

其中序遍历的结果是 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 的左子树。

现在有两个问题：

A 放哪里？
根节点 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
}

结尾

上文使用了二叉排序树作为辅助理解的工具，顺便一提，在想到二叉排序树的时候，还应将其和二分搜索联系起来。

单看树的旋转，其实非常简单。树的旋转实际上是为平衡二叉树做铺垫，因此下一篇将会把普通的二叉树换成平衡二叉树。

Tags: 算法与数据结构