题目描述

在二叉树中增加一行
给定一个二叉树，根节点为第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]。
输入的二叉树至少有一个节点。

算法1

时间复杂度O(n)

Java 代码

class Solution {
    public TreeNode addOneRow(TreeNode root, int v, int d) {
        if (root == null) {
            if (d == 1) {
                return new TreeNode(v);
            } else {
                return null;
            }
        }

        if (d == 1) {
            TreeNode t = new TreeNode(v);
            t.left = root;
            return t;
        }

        Queue<TreeNode> q = new ArrayDeque<>();
        q.offer(root);
        int depth = 1;
        // for now, d=[2,depth+1]
        while (!q.isEmpty()) {
            int levelSize = q.size();
            for (int i = 0; i < levelSize; i++) {
                TreeNode t = q.poll();
                if (depth == d - 1) {
                    TreeNode left = new TreeNode(v);
                    TreeNode right = new TreeNode(v);
                    if (t.left != null) {
                        left.left = t.left;
                    }
                    if (t.right != null) {
                        right.right = t.right;
                    }
                    t.left = left;
                    t.right = right;
                    // 重要continue!
                    continue;
                }
                if (t.left != null) {
                    q.offer(t.left);
                }

                if (t.right != null) {
                    q.offer(t.right);
                }
            }

            depth++;
        }

        return root;
    }
}