在二叉树中增加一行

难度: Medium

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

提示:

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

Submission

运行时间: 27 ms

内存: 17.5 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[TreeNode]:
        if not root:
            return TreeNode(val)
        if depth == 1:
            newRoot = TreeNode(val)
            newRoot.left = root
            return newRoot
        self.dp(root,1,depth-1,val)
        return root
        
    def dp(self,root,depth,target_depth,val):
        if not root:
            return 
        if depth == target_depth:
            newleft  = TreeNode(val)
            newright = TreeNode(val)
            newleft.left = root.left
            newright.right = root.right
            root.left = newleft
            root.right = newright
        depth += 1
        self.dp(root.left, depth, target_depth, val)
        self.dp(root.right, depth, target_depth, val)
        depth -= 1
        return root

Explain

这个题解采用递归的方式来解决问题。主要思路是通过递归遍历二叉树，当到达指定深度 depth-1 时，创建两个新节点，将原有的左右子树挂载到新创建的节点上，然后将新节点作为当前节点的新的左右子节点。如果 depth 为1，则直接创建一个新节点作为根节点，原有的整棵树作为新根节点的左子树。

时间复杂度: O(n)

空间复杂度: O(h)，其中 h 为二叉树的高度，h ∈ [log(n), n]

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def addOneRow(self, root: Optional[TreeNode], val: int, depth: int) -> Optional[TreeNode]:
        # 如果根节点为空，直接返回一个新节点
        if not root:
            return TreeNode(val)
        # 如果指定深度为1，创建一个新的根节点，原树作为新根的左子树
        if depth == 1:
            newRoot = TreeNode(val)
            newRoot.left = root
            return newRoot
        # 递归遍历树，在指定深度处添加新行
        self.dp(root, 1, depth-1, val)
        return root
        
    def dp(self, root, depth, target_depth, val):
        # 如果当前节点为空，直接返回
        if not root:
            return 
        # 如果当前深度等于目标深度，创建新的左右子节点
        if depth == target_depth:
            newleft  = TreeNode(val)
            newright = TreeNode(val)
            newleft.left = root.left
            newright.right = root.right
            root.left = newleft
            root.right = newright
        # 递归遍历左右子树
        depth += 1
        self.dp(root.left, depth, target_depth, val)
        self.dp(root.right, depth, target_depth, val)
        depth -= 1
        return root

Explore

这样做主要是为了确保当递归到目标深度时，可以直接在当前节点插入新的节点，然后再递归遍历原有的子树。如果不这样做，即在递归遍历子树后再添加新节点，那么新添加的节点将不会被进一步递归处理，导致新节点下的子树结构不完整。因此，这种顺序确保了每个新插入的节点都正确地连接了其应有的子树。

确实，当前的实现在达到目标深度之后继续对所有子节点进行递归调用，这在技术上是不必要的。优化方法是在达到目标深度并完成节点插入后，停止进一步的递归调用。这样可以减少递归的深度，提高算法的效率。具体来说，可以在插入新节点后直接返回，而不是继续递归调用。

在递归过程中，每次递归调用都会通过参数传递当前的深度值。在函数内部，深度值`depth`在每次递归调用前被增加1，调用结束后，由于每次调用都有自己的局部变量`depth`，这个增加的操作不会影响到其他递归调用中的`depth`值。因此，每个递归调用都独立维护了自己的深度计数，避免了对共享状态的修改，确保了变量值的正确性和函数的可重入性。

当`depth`为1时，题解中创建了一个新的根节点，并将整个原始树设置为新根节点的左子树。原始树的右子树不会发生变化，但是它将继续作为原始根节点的右子树存在。新根节点的右子树在这种情况下是空的，因为题解并未提及将任何部分设置为新根的右子树。