LC 623. 在二叉树中增加一行

题目描述

这是 LeetCode 上的 623. 在二叉树中增加一行 ，难度为中等。

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

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

加法规则如下:

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

示例 1:

1
2
3

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

输出: [4,1,1,2,null,null,6,3,1,5]

示例 2:

1
2
3

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

输出:  [4,2,null,1,1,3,null,null,1]

提示:

节点数在 $[1, 10^4]$ 范围内
树的深度在 $[1, 10^4]$范围内
$-100 <= Node.val <= 100$
$-10^5 <= val <= 10^5$
$1 <= depth <= the depth of tree + 1$

BFS

根据 BFS 来做，每次 BFS 将整一层进行拓展，同时记录当前深度，当到达第 depth - 1 层，则进行加点操作。

Java 代码：

class Solution {
    public TreeNode addOneRow(TreeNode root, int val, int depth) {
        if (depth == 1) return new TreeNode(val, root, null);
        Deque<TreeNode> d = new ArrayDeque<>();
        d.addLast(root);
        int cur = 1;
        while (!d.isEmpty()) {
            int sz = d.size();
            while (sz-- > 0) {
                TreeNode t = d.pollFirst();
                if (cur == depth - 1) {
                    TreeNode a = new TreeNode(val), b = new TreeNode(val);
                    a.left = t.left; b.right = t.right;
                    t.left = a; t.right = b;
                } else {
                    if (t.left != null) d.addLast(t.left);
                    if (t.right != null) d.addLast(t.right);
                }
            }
            cur++;
        }
        return root;
    }
}

C++ 代码：

class Solution {
public:
    TreeNode* addOneRow(TreeNode* root, int val, int depth) {
       if (depth == 1) return new TreeNode(val, root, nullptr);
        queue<TreeNode*> q;
        q.push(root);
        int cur = 1;
        while (!q.empty()) {
            int sz = q.size();
            while (sz--) {
                TreeNode* t = q.front();
                q.pop();
                if (cur == depth - 1) {
                    TreeNode* a = new TreeNode(val);
                    TreeNode* b = new TreeNode(val);
                    a->left = t->left; b->right = t->right;
                    t->left = a; t->right = b;
                } else {
                    if (t->left) q.push(t->left);
                    if (t->right) q.push(t->right);
                }
            }
            cur++;
        }
        return root;
    }
};

Python 代码：

class Solution:
    def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[TreeNode]:
        if depth == 1:
            return TreeNode(val, root, None)
        d = deque([root])
        cur = 1
        while d:
            sz = len(d)
            while sz > 0:
                t = d.popleft()
                if cur == depth - 1:
                    a, b = TreeNode(val), TreeNode(val)
                    a.left, b.right = t.left, t.right
                    t.left, t.right = a, b
                else:
                    if t.left:
                        d.append(t.left)
                    if t.right:
                        d.append(t.right)
                sz -= 1
            cur += 1
        return root

TypeScript 代码：

function addOneRow(root: TreeNode | null, val: number, depth: number): TreeNode | null {
    if (depth == 1) return new TreeNode(v, root, null);
    const stk = new Array<TreeNode>()
    let he = 0, ta = 0, cur = 1
    stk[ta++] = root
    while (he < ta) {
        let sz = ta - he
        while (sz-- > 0) {
            const t = stk[he++]
            if (cur == depth - 1) {
                const a = new TreeNode(val), b = new TreeNode(val)
                a.left = t.left; b.right = t.right
                t.left = a; t.right = b
            } else {
                if (t.left != null) stk[ta++] = t.left
                if (t.right != null) stk[ta++] = t.right
            }
        }
        cur++
    }
    return root
};

时间复杂度：$O(n)$
空间复杂度：$O(n)$

DFS

同理，使用 DFS 也可进行求解，在 DFS 过程中记录当前深度。

Java 代码：

class Solution {
    int d, v;
    public TreeNode addOneRow(TreeNode root, int val, int depth) {
        d = depth; v = val;
        if (d == 1) return new TreeNode(val, root, null);
        dfs(root, 1);
        return root;
    }
    void dfs(TreeNode root, int cur) {
        if (root == null) return ;
        if (cur == d - 1) {
            TreeNode a = new TreeNode(v), b = new TreeNode(v);
            a.left = root.left; b.right = root.right;
            root.left = a; root.right = b;
        } else {
            dfs(root.left, cur + 1);
            dfs(root.right, cur + 1);
        }
    }
}

C++ 代码：

class Solution {
public:
    int d, v;
    TreeNode* addOneRow(TreeNode* root, int val, int depth) {
        d = depth; v = val;
        if (d == 1) return new TreeNode(val, root, nullptr);
        dfs(root, 1);
        return root;
    }
    void dfs(TreeNode* root, int cur) {
        if (root == nullptr) return;
        if (cur == d - 1) {
            TreeNode* a = new TreeNode(v);
            TreeNode* b = new TreeNode(v);
            a->left = root->left; b->right = root->right;
            root->left = a; root->right = b;
        } else {
            dfs(root->left, cur + 1);
            dfs(root->right, cur + 1);
        }
    }
};

Python 代码：

class Solution:
    def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[TreeNode]:
        self.d = depth
        self.v = val
        if depth == 1:
            return TreeNode(val, root, None)
        self.dfs(root, 1)
        return root
    
    def dfs(self, root, cur):
        if not root:
            return
        if cur == self.d - 1:
            a, b = TreeNode(self.v), TreeNode(self.v)
            a.left, b.right = root.left, root.right
            root.left, root.right = a, b
        else:
            self.dfs(root.left, cur + 1)
            self.dfs(root.right, cur + 1)

TypeScript 代码：

let d = 0, v = 0
function addOneRow(root: TreeNode | null, val: number, depth: number): TreeNode | null {
    d = depth; v = val
    if (d == 1) return new TreeNode(v, root, null);
    dfs(root, 1)
    return root
};
function dfs(root: TreeNode | null, cur: number): void {
    if (root == null) return 
    if (cur == d - 1) {
        const a = new TreeNode(v), b = new TreeNode(v)
        a.left = root.left; b.right = root.right
        root.left = a; root.right = b
    } else {
        dfs(root.left, cur + 1)
        dfs(root.right, cur + 1)
    }
}

时间复杂度：$O(n)$
空间复杂度：$O(n)$

最后

这是我们「刷穿 LeetCode」系列文章的第 No.623 篇，系列开始于 2021/01/01，截止于起始日 LeetCode 上共有 1916 道题目，部分是有锁题，我们将先把所有不带锁的题目刷完。

在这个系列文章里面，除了讲解解题思路以外，还会尽可能给出最为简洁的代码。如果涉及通解还会相应的代码模板。

为了方便各位同学能够电脑上进行调试和提交代码，我建立了相关的仓库：https://github.com/SharingSource/LogicStack-LeetCode 。

在仓库地址里，你可以看到系列文章的题解链接、系列文章的相应代码、LeetCode 原题链接和其他优选题解。

DFS BFS 二叉树

本博客所有文章除特别声明外，均采用 CC BY-SA 4.0 协议，转载请注明出处！

LC 1408. 数组中的字符串匹配上一篇

LC 剑指 Offer 34. 二叉树中和为某一值的路径下一篇