623_add_one_row_to_tree

623. 在二叉树中增加一行 中等

给定一个二叉树，根节点为第1层，深度为 1。在其第 d 层追加一行值为 v 的节点。

添加规则：给定一个深度值 d （正整数），针对深度为 d-1 层的每一非空节点 N，为 N 创建两个值为 v 的左子树和右子树。

将 N 原先的左子树，连接为新节点 v 的左子树；将 N 原先的右子树，连接为新节点 v 的右子树。

如果 d 的值为 1，深度 d - 1 不存在，则创建一个新的根节点 v，原先的整棵树将作为 v 的左子树。

示例 1:

输入: 
二叉树如下所示:
       4
     /   \
    2     6
   / \   / 
  3   1 5   

v = 1

d = 2

输出: 
       4
      / \
     1   1
    /     \
   2       6
  / \     / 
 3   1   5   

示例 2:

输入: 
二叉树如下所示:
      4
     /   
    2    
   / \   
  3   1    

v = 1

d = 3

输出: 
      4
     /   
    2
   / \    
  1   1
 /     \  
3       1

注意:

• 输入的深度值 d 的范围是：[1，二叉树最大深度 + 1]。
• 输入的二叉树至少有一个节点。

package main

import "fmt"

func main() {
    root := &TreeNode{
        Val:   4,
        Left:  &TreeNode{Val: 2},
        Right: &TreeNode{Val: 6},
    }

    newRoot := addOneRow(root, 1, 2)
    fmt.Println(newRoot.Val, newRoot.Left.Val, newRoot.Right.Val) // 4 1 1
    // also ok
}

// 层序遍历即可
func addOneRow(root *TreeNode, v int, d int) *TreeNode {
    if root == nil {
        return nil
    }
    if d == 1 {
        return &TreeNode{Val: v, Left: root} // 特殊 case
    }

    depth := 1
    q := []*TreeNode{root}
    for len(q) > 0 {
        counts := len(q)
        for i := 0; i < counts; i++ {
            cur := q[0]
            q = q[1:]
            if depth == d-1 { // 注意 -1
                cur.Left = &TreeNode{Val: v, Left: cur.Left}
                cur.Right = &TreeNode{Val: v, Right: cur.Right}
            }
            if cur.Left != nil {
                q = append(q, cur.Left)
            }
            if cur.Right != nil {
                q = append(q, cur.Right)
            }
        }
        depth++
    }
    return root
}