[LC623]在二叉树中增加一行

题目描述

给定一个二叉树的根 root 和两个整数 val 和 depth ，在给定的深度 depth 处添加一个值为 val 的节点行。

注意，根节点 root 位于深度 1 。

加法规则如下:

给定整数 depth，对于深度为 depth - 1 的每个非空树节点 cur ，创建两个值为 val 的树节点作为 cur 的左子树根和右子树根。
cur 原来的左子树应该是新的左子树根的左子树。
cur 原来的右子树应该是新的右子树根的右子树。
如果 depth == 1 意味着 depth - 1 根本没有深度，那么创建一个树节点，值 val 作为整个原始树的新根，而原始树就是新根的左子树。

来源：力扣（LeetCode）
链接：https://leetcode.cn/problems/add-one-row-to-tree

题目解析

一道经典的树的遍历题。通常树的遍历题可以用深度优先遍历与广度遍历解决，本题同样适用。对于这样，需要在特定深度进行统一操作的，广度优先可能是更为自然的做法。

代码

BFS

func addOneRow(root *TreeNode, val int, depth int) *TreeNode {
	if depth == 1 {
		return &TreeNode{val, root, nil}
	}
	queue := []*TreeNode{root}
	//  locate the nodes at depth - 1
	for i := 1; i < depth-1; i++ {
		tempQueue := queue
		queue = nil
		for _, node := range tempQueue {
			if node.Left != nil {
				queue = append(queue, node.Left)
			}
			if node.Right != nil {
				queue = append(queue, node.Right)
			}
		}
	}
	// create new nodes at depth 
	for _, node := range queue {
		node.Left = &TreeNode{val, node.Left, nil}
		node.Right = &TreeNode{val, nil, node.Right}
	}
	return root
}

DFS

func addOneRow(root *TreeNode, val int, depth int) *TreeNode {
	if root == nil {
		return root
	}
	if depth == 1 {
		return &TreeNode{val, root, nil}
	}
	if depth == 2 {
		root.Left = &TreeNode{val, root.Left, nil}
		root.Right = &TreeNode{val, nil, root.Right}
		return root
	}
	root.Left = addOneRow(root.Left, val, depth-1)
	root.Right = addOneRow(root.Right, val, depth-1)
	return root
}