【树】在二叉树中增加一行

问题描述

给定一个二叉树的根 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]

解题思路

采用层序遍历，每遍历一层，校验depth是否满足条件，如果满足条件则增加节点做计算。

示例1中流程如下：

遍历第1层，只有节点4

由于depth要求是第2层，则需要对第1层数据作为根节点，第2层作为叶子节点

-- 执行逻辑，新建节点Node，第1层的节点left指向新建节点Node，将新建节点Node指向原来left指向的左子树。

-- 执行逻辑，新建节点Node，第1层的节点right指向新建节点Node，将新建节点Node指向原来right指向的右子树。

代码实现

import org.example.TreeNode;

import java.util.Deque;
import java.util.LinkedList;

class Solution {
    public TreeNode addOneRow(TreeNode root, int val, int depth) {

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

        TreeNode root0 = root;
        Deque<TreeNode> deque = new LinkedList<>();
        if (root != null) {
            deque.offer(root);
        }

        int count = 0;
        while (!deque.isEmpty()) {
            count++;
            int sz = deque.size();
            for (int i = 0; i < sz; i++) {
                TreeNode node = deque.poll();

                if (count == depth - 1) {
                    TreeNode left = node.left;
                    TreeNode right = node.right;

                    node.left = new TreeNode(val);
                    node.left.left = left;

                    node.right = new TreeNode(val);
                    node.right.right = right;

                } else {
                    if (node.left != null) {
                        deque.offer(node.left);
                    }
                    if (node.right != null) {
                        deque.offer(node.right);
                    }
                }
            }

        }
        return root0;
    }
}

二叉树TreeNode的实现逻辑：

package org.example;

public class TreeNode {
    public int val;
    public TreeNode left;
    public TreeNode right;

    public TreeNode() {
    }

    public TreeNode(int val) {
        this.val = val;
    }

    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}