We could use BFS to solve this problem. In the queue of BFS, we save one additional information which is depth of this node.
# Definition for a binary tree node.
# class TreeNode:
# def __init__(self, x):
# self.val = x
# self.left = None
# self.right = None
from collections import deque
class Solution:
def addOneRow(self, root: TreeNode, v: int, d: int) -> TreeNode:
if d == 1:
ans = TreeNode(v)
ans.left = root
return ans
queue = deque([])
queue.append([root, 0])
while queue and queue[0][1] < d - 2:
node, depth = queue.popleft()
if node.left:
queue.append([node.left, depth + 1])
if node.right:
queue.append([node.right, depth + 1])
while queue:
node, _ = queue.pop()
l = TreeNode(v)
r = TreeNode(v)
l.left = node.left
r.right = node.right
node.left, node.right = l, r
return root