Add One Row to Tree Python Solution with Tests

Tested Python solution for LeetCode 623 with 21 pytest cases. Generate a practice environment with lcpy.

LeetCode 623, Medium. Topics: Tree, Depth-First Search, Breadth-First Search, Binary Tree. View on LeetCode.

Generate this problem as a practice environment: tested reference solution, 21 parametrized pytest cases, and a playground notebook:

lcpy gen -n 623   # by problem number
lcpy gen -s add_one_row_to_tree   # by problem name

​

Problem

Given the root of a binary tree and two integers val and depth, add a row of nodes with value val at the given depth depth.

Note that the root node is at depth 1.

The adding rule is:

• Given the integer depth, for each not null tree node cur at the depth depth - 1, create two tree nodes with value val as cur’s left subtree root and right subtree root.
• cur’s original left subtree should be the left subtree of the new left subtree root.
• cur’s original right subtree should be the right subtree of the new right subtree root.
• If depth == 1 that means there is no depth depth - 1 at all, then create a tree node with value val as the new root of the whole original tree, and the original tree is the new root’s left subtree.

​

Examples

Input: root = [4,2,6,3,1,5], val = 1, depth = 2
Output: [4,1,1,2,null,null,6,3,1,5]

Input: root = [4,2,null,3,1], val = 1, depth = 3
Output: [4,2,null,1,1,3,null,null,1]

​

Constraints

• The number of nodes in the tree is in the range [1, 10^4]
• The depth of the tree is in the range [1, 10^4]
• -100 <= Node.val <= 100
• -10^5 <= val <= 10^5
• 1 <= depth <= the depth of tree + 1

​

Solution

Reference implementation from solution.py on GitHub, full suite in test_solution.py:

from collections import deque

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(w) where w is the maximum width of the tree
    def add_one_row(self, root: TreeNode[int] | None, val: int, depth: int) -> TreeNode[int] | None:
        if depth == 1:
            new_root = TreeNode(val)
            new_root.left = root
            return new_root

        queue: deque[TreeNode[int]] = deque()
        if root is not None:
            queue.append(root)

        current_depth = 1
        while queue and current_depth < depth - 1:
            for _ in range(len(queue)):
                node = queue.popleft()
                if node.left is not None:
                    queue.append(node.left)
                if node.right is not None:
                    queue.append(node.right)
            current_depth += 1

        for parent in queue:
            left = TreeNode(val)
            left.left = parent.left
            parent.left = left

            right = TreeNode(val)
            right.right = parent.right
            parent.right = right

        return root

​

Complexity

Time	Space
O(n)	O(w) where w is the maximum width of the tree

​

Tags