LeetCode 每日一题——623. 在二叉树中增加一行

1.题目描述

623. 在二叉树中增加一行

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

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

加法规则如下:

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

示例 1:

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

示例 2:

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

2.解题思路与代码

2.1 解题思路

这道题因为需要在指定深度的层上加一层，所以考虑使用 BFS 进行解答。首先考虑当给定 depth 等于 1 的时候，即要把 val 插入到根节点之上，按照题目要求直接创建新的根节点 newRoot，并且让新的根节点的 left 指向原来的根节点，即 newRoot.left = root，代码如下：

if (depth == 1) {
    TreeNode newRoot = new TreeNode(val);
    newRoot.left = root;
    return newRoot;
}

接下来便开始处理 depth 大于 1 的情况，我们还是套用二叉树 BFS 代码公式进行解答。使用一个队列存储遍历的节点，并使用变量 layer 和 count 分别统计遍历的层数和放入队列的次数。首先进行初始化将根节点放入队列中，因为根节点入队因此此时 count 为 1，并且根据题目要求 layer 层数也是 1 。当放入队列次数等于队列大小的时候，表示一层已经遍历完了，此时层数加 1，并且让 count 归零，代码如下：

if (queue.size() == count) {
    layer++;
    count = 0;
}

这个时候我们从队列中拿出第一个节点，此时判断层数 layer 是否等于 depth，如果不相等则将弹出节点的左右节点入栈；如果相等，那么就需要使用 val 新建两个节点，然后让弹出节点的左右子树指向新建节点，并且按照题目要求原来的左子树应该是新的左子树根的左子树，右子树应该是新的右子树根的右子树，代码如下：

if (layer == depth) {
    TreeNode leftNode = new TreeNode(val);
    TreeNode rightNode = new TreeNode(val);
    leftNode.left = poll.left;
    poll.left = leftNode;
    rightNode.right = poll.right;
    poll.right = rightNode;
}

需要注意的是我们需要在二叉树的一层插入，所以上面的插入执行完成之后不能够退出循环，因为不能够保证这一层是否遍历完成，只有当层数 layer 大于 depth 的时候我们便可以退出循环结果返回根节点。

2.2 代码

class Solution {
    public TreeNode addOneRow(TreeNode root, int val, int depth) {
        // 在根节点插入上一层
        if (depth == 1) {
            TreeNode newRoot = new TreeNode(val);
            newRoot.left = root;
            return newRoot;
        }
        LinkedList<TreeNode> queue = new LinkedList<>();
        int layer = 1;
        int count = 1;
        queue.offer(root);
        while (!queue.isEmpty()) {
            // queue.size() == count 时表示一层已经遍历完成，遍历完一层时层数加一，并且count归零
            if (queue.size() == count) {
                layer++;
                count = 0;
            }
            TreeNode poll = queue.poll();
            // 插入新节点
            if (layer == depth) {
                TreeNode leftNode = new TreeNode(val);
                TreeNode rightNode = new TreeNode(val);
                leftNode.left = poll.left;
                poll.left = leftNode;
                rightNode.right = poll.right;
                poll.right = rightNode;
            }
            // 层数大于 depth 时退出循环
            if (layer > depth) {
                break;
            }
            if (poll.left != null) {
                queue.offer(poll.left);
                count++;
            }
            if (poll.right != null) {
                queue.offer(poll.right);
                count++;
            }
        }
        return root;
    }
}

2.3 测试结果

通过测试

3.总结

• 使用二叉树 BFS 进行解答