Add One Row to Tree

Problem Id: 623 Difficulty: Medium Tag: Tree


Intuition

We could use BFS to solve this problem. In the queue of BFS, we save one additional information which is depth of this node.

Solution


# 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