LC623. Add One Row to Tree

# 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]:
        target_depth = depth
 
        def dfs(node, depth = 1):
            if not node: return None
 
            # Change left and right as we reach parent depth nodes
            if depth == target_depth - 1:
                node.left = TreeNode(val, node.left, None)
                node.right = TreeNode(val, None, node.right)
                return node
            
            dfs(node.left, depth + 1)
            dfs(node.right, depth + 1)
            
            return node
        
        if target_depth == 1:
            return TreeNode(val, root, None)
        
        return dfs(root)

H/T @amol1729 for helping simplify the original approach.

Complexity

Time: $O(n)$
Space: $O(n)$