背景

之前在瞭解二叉樹旋轉的時候，為了方便檢視中間狀態，就寫了個以樹狀形式列印二叉樹的函式。

起初是使用二叉樹中序遍歷的結果展開的方式，簡單但打印出來的樹有一定的傾斜。

例如這棵樹：

  5
 3  7
2  6 8

它的中序遍歷結果為：

+++++++++++++
|2|3|5|6|7|8|
+++++++++++++

打印出來的結果中，節點 3 和節點 7 不是對稱的。因為節點 3 距離其父節點 5 的距離只有 1，而節點 7 距離其父節點 5 的距離則是 2。

於是做了一番改造，列印了對稱的樹：

   5
 3   7
2   6 8

對應的陣列是：

+++++++++++++++
|2|3| |5|6|7|8|
+++++++++++++++
 

中序遍歷

這裡改成返回遍歷結果而不是直接列印。

// InorderIteration 中序遍歷迭代法
func InorderIteration(root *TreeNode) []*TreeNode {
	rs := make([]*TreeNode, 0)
	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]

		rs = append(rs, root)

		root = root.Right
	}

	return rs
}
 

基於中序遍歷結果的展開

// PrintTree 以中序遍歷結果展開樹
func PrintTree(root *TreeNode) {
	if root == nil {
		return
	}
	inorder := InorderIteration(root)

	row := []*TreeNode{root}
	var cache []*TreeNode
	for len(row) != 0 {
		// 每次取一整行出來，並重新申請空間存放下一行
		cache = row
		row = make([]*TreeNode, 0)

		// 從中序遍歷中尋找當前行的資料，不匹配的列印空格
		i := 0
		for _, node := range inorder {
			if node.Value != cache[i].Value {
				fmt.Print(" ")
			} else {
				fmt.Print(node.Value)
				i++
				if i == len(cache) {
					break
				}
			}
		}
		fmt.Println()

		for _, node := range cache {
			if node.Left != nil {
				row = append(row, node.Left)
			}
			if node.Right != nil {
				row = append(row, node.Right)
			}
		}
	}
}
 

補全空位置的列印

借鑑中序遍歷展開的思路。根據樹的高度，申請一個可以容納這個高度滿節點狀態的節點數量的陣列。從根節點開始，寬度優先遍歷。讓每個節點都把一個範圍平均分成兩部分。

為了方便，這裡先展示具有破壞性的列印。破壞的地方為 Node 的 Height 屬性。在列印時，會發生變化。如果不想破壞，則再增加一個 Layer 屬性即可。

func PrintTree(root *TreeNode) {
	// 把子節點的高度更新為父節點高度-1
	UpdateHightToMax(root)
	data := ToStrictInorderArray(root)
	height := root.Height
	for height > 0 {
		// 每次遍歷完整的中序遍歷結果，當元素的層級與當前層級相同時列印值，不同時列印空格
		for _, element := range data {
			if element == nil || element.Height != height {
				fmt.Print(" ")
			} else {
				fmt.Printf("%d", element.Value)
			}
		}
		height--
		fmt.Println()
	}
}

// UpdateHightToMax 前序遍歷更新高度，把子節點的高度更新為父節點高度-1
func UpdateHightToMax(root *TreeNode) {
	stack := make([]*TreeNode, 0)
	for len(stack) != 0 || root != nil {
		for root != nil {
			if root.Parent != nil {
				root.Height = root.Parent.Height - 1
			}
			stack = append(stack, root)
			root = root.Left
		}

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

		root = node.Right
	}
}

// ToStrictInorderArray 對稱的樹
func ToStrictInorderArray(root *TreeNode) []*TreeNode {
	if root == nil || root.Height < 0 {
		return make([]*TreeNode, 0)
	}

	// 總節點數為 (2^n)-1
	result := make([]*TreeNode, (1<<root.Height)-1)

	// 確保操作下一層節點時，父節點已放入結果集
	queue := BreadthFirstIteration(root)

	// root 的位置無法從父節點算出，先算出來
	index := 1<<(root.Height-1) - 1
	result[index] = queue[0]
	queue = queue[1:]

	for len(queue) > 0 {
		iterNode := queue[0]
		queue = queue[1:]

		distance := 1 << (iterNode.Height - 1)

		idx := FindIndexStrictInorder(result, iterNode.Parent)
		if iterNode.Value < iterNode.Parent.Value {
			idx -= distance
		} else {
			idx += distance
		}

		result[idx] = iterNode
	}

	return result
}

// BreadthFirstIteration 廣度優先遍歷
func BreadthFirstIteration(node *TreeNode) []*TreeNode {
	data := []*TreeNode{node}
	for index := 0; index < len(data); index++ {
		node = data[index]
		if node.Left != nil {
			data = append(data, node.Left)
		}
		if node.Right != nil {
			data = append(data, node.Right)
		}
	}
	return data
}

// FindIndexStrictInorder 在陣列中找到指定節點的位置
func FindIndexStrictInorder(tree []*TreeNode, node *TreeNode) int {
	cut := (len(tree) + 1) >> 1
	middle := cut - 1
	for cut > 1 && middle <= len(tree) {
		cut = cut >> 1
		if node.Value < tree[middle].Value {
			middle -= cut
		} else if node.Value > tree[middle].Value {
			middle += cut
		} else {
			return middle
		}
	}

	return middle
}