携手创作，共同成长！这是我参与「掘金日新计划 · 8 月更文挑战」的第10天，点击查看活动详情

刷题的日常

一天一题，保持脑子清爽

在二叉树中增加一行

来自leetcode的 623 题，题意如下：

给定一个二叉树的根root和两个整数 val 和depth，在给定的深度depth处添加一个值为 val 的节点行。
注意，根节点root位于深度1。
加法规则如下:

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

示例如下：

输入: root = [4,2,6,3,1,5], val = 1, depth = 2
输出: [4,1,1,2,null,null,6,3,1,5]

理解题意

我们可以从题意中提取的条件如下：

• 给定一颗树，并给定一个值和树的层数
• 要求在给定的层数插入新的一层
• 新一层的树节点的值为给定的值
• 要注意层数为一的情况，这个时候根节点会被替换成新的节点
• 不是根节点的话，左子树放新节点的左边，右子树放在新节点的右边

做题思路

首先做个特殊判断，如果是在根节点做的替换，我们需要生成新的根节点，然后将原根节点放在新节点的左边并返回。
其它情况需要先找到要插入层数的父级：

• 对树进行广度便利
• 当前层数如果是父级就找到了需要插入的地方，跳出循环
• 队列里保存了当前需要插入的父节点信息
• 便利当前的队列，按照规则插入新的子节点
• 返回原来的根节点

代码实现

代码实现如下：

public class Solution {
    public TreeNode addOneRow(TreeNode root, int val, int depth) {
        if (depth == 1) {
            TreeNode result = new TreeNode(val);
            result.left = root;
            return result;
        }
        int size;
        Queue<TreeNode> queue = new LinkedList<>();
        queue.add(root);

        while (--depth > 1) {
            size = queue.size();
            while (size-- > 0) {
                TreeNode node = queue.poll();
                if (node.left != null) {
                    queue.add(node.left);
                }
                if (node.right != null) {
                    queue.add(node.right);
                }
            }
        }
        TreeNode left, right;
        TreeNode newLeft, newRight;
        size = queue.size();
        while (size-- > 0) {
            TreeNode node = queue.poll();
            left = node.left;
            newLeft = new TreeNode(val);
            newLeft.left = left;
            node.left = newLeft;
                
            right = node.right;
            newRight = new TreeNode(val);
            newRight.right = right;
            node.right = newRight;
        }
        return root;
    }
}