题目来源牛客练习题
层序遍历通过队列来实现,每次将当前节点的子节点添加到队列尾部,再从队列头部取出一个元素即可。
此道题目要求按层分成不同的数组,因此多一个步骤,要记录该节点属于的层数。
添加一个结构体即可,当队列为空或者当前节点和下一个节点的层数不一样,就将累加的数组添入答案即可。
代码如下所示:
type data struct {
node *TreeNode
floor int
}
func levelOrder(root *TreeNode) [][]int {
// write code here
ans := [][]int{}
if root == nil {
return ans
}
stack := []data{}
stack = append(stack, data{root, 0})
tmp := []int{}
for len(stack) != 0 {
//取出头部
nowNode := stack[0]
stack = stack[1:]
//尾部添加
if nowNode.node.Left != nil {
stack = append(stack, data{nowNode.node.Left, nowNode.floor + 1})
}
if nowNode.node.Right != nil {
stack = append(stack, data{nowNode.node.Right, nowNode.floor + 1})
}
//将当前节点Val添入tmp
tmp = append(tmp, nowNode.node.Val)
if len(stack) == 0 || stack[0].floor != nowNode.floor {
//因为使用的是同一个数组,地址相同,直接添加可能会后面的更新前面的值
//因此需要每次声明一个局部变量,添加进去
tmpData:=[]int{}
tmpData=append(tmpData, tmp...)
ans = append(ans, tmpData)
//tmp添加进答案后重置
tmp = tmp[:0]
}
}
return ans
}