题目链接

给定一个二叉树，根节点为第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

注意:

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

解答

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode() {}
 *     TreeNode(int val) { this.val = val; }
 *     TreeNode(int val, TreeNode left, TreeNode right) {
 *         this.val = val;
 *         this.left = left;
 *         this.right = right;
 *     }
 * }
 */
class Solution {
    public TreeNode addOneRow(TreeNode root, int val, int depth) {
        if (depth == 1) return new TreeNode(val, root, null);
        TreeNode node;
        int currLevelNodeNum = 1, nextLevelNodeNum = 0, currDepth = 1;
        Queue<TreeNode> nodeQueue = new LinkedList<>();
        nodeQueue.add(root);
        while (!nodeQueue.isEmpty()) {
            node = nodeQueue.poll();
            if (currDepth == depth-1) {
                // 当前层即为要插入一行的元素的上一行
                node.left = new TreeNode(val, node.left, null);
                node.right = new TreeNode(val, null, node.right);
            }
            if (currDepth < depth) {
                if (node.left != null) { nodeQueue.add(node.left); nextLevelNodeNum++; }
                if (node.right != null) { nodeQueue.add(node.right); nextLevelNodeNum++; }
            }
            if (--currLevelNodeNum == 0) {
                if (currDepth == depth-1) break;
                currLevelNodeNum = nextLevelNodeNum;
                nextLevelNodeNum = 0;
                currDepth++;
            }
        }
        return root;
    }
}