https://leetcode.com/problems/add-one-row-to-tree

The idea here is to use recursive dfs to build our new tree. Let us use dfs(node, h, dr) function, where

  1. node is node we are currently traversing.
  2. h is height of node we are currently on.
  3. dr is direction from which we come to this node, it is 0 if node is left children of another node and 1 if node is right children of another node.

When we traverse tree, using dfs, we need to check if h == d, that is if current height of node is equal to row we need to insert. In this case we create new node tmp with value v and insert it: if dr = 0, we need to attach node as left children of tmp and if dr = 1, we need to attach it as right children. Also there is case, if node equal to None (during traversal we allow our dfs visit None nodes), then we do not need to attach anything to it, but in fact this case will be covered by previous two cases.

If h != d, then we need to run our function recursively for both children and return node in the end.

Complexity: time complexity is O(n): we traverse our tree once and add not more than n nodes. Space complexity is O(h) - height of tree if we do not count answer and O(h + n+ q), where q is number of nodes we added if we count anwer space.

class Solution:
    def addOneRow(self, root, v, d):
        def dfs(node, h, dr):
            
            if h == d:
                tmp = TreeNode(v)
                #if not node: return tmp
                if dr == 0:
                    tmp.left = node
                else:
                    tmp.right = node
                return tmp
            
            if not node: return node
            
            node.left = dfs(node.left, h+1, 0)
            node.right = dfs(node.right, h+1, 1)
            return node
            
        return dfs(root, 1, 0)

If you like the solution, you can upvote it on leetcode discussion section: Problem 0623