二叉树的层序遍历

原二叉树:
   3
   / \
  9  20
    /  \
   15   7
要求遍历结果:
[
  [3],
  [9,20],
  [15,7]
]

代码及注释

package main

import "fmt"

func main() {
	root := buildTree()
	fmt.Println(levelOrder(root))
}

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

func levelOrder(root *TreeNode) [][]int {
	if root == nil {
		return nil
	}
	res := [][]int{}
	cLevelNodes := []*TreeNode{root}        // 当前层节点
	for i := 0; len(cLevelNodes) > 0; i++ { // 每次遍历当前层节点
		res = append(res, []int{})
		// 定义下一层节点数据
		nLevelNodes := []*TreeNode{}
		for j := 0; j < len(cLevelNodes); j++ { // 内层循环将当前层的所有节点取出
			res[i] = append(res[i], cLevelNodes[j].Val)
			// 获取当前层的左右节点，并赋值给nLevelNodes
			if cLevelNodes[j].Left != nil {
				nLevelNodes = append(nLevelNodes, cLevelNodes[j].Left)
			}
			if cLevelNodes[j].Right != nil {
				nLevelNodes = append(nLevelNodes, cLevelNodes[j].Right)
			}
		}
		cLevelNodes = nLevelNodes // 继续遍历下一层
	}
	return res
}

func buildTree() *TreeNode {
	return &TreeNode{
		Val: 3,
		Left: &TreeNode{
			Val:   9,
			Left:  nil,
			Right: nil,
		},
		Right: &TreeNode{
			Val: 20,
			Left: &TreeNode{
				Val:   15,
				Left:  nil,
				Right: nil,
			},
			Right: &TreeNode{
				Val:   7,
				Left:  nil,
				Right: nil,
			},
		},
	}
}