leetcode - 1110. Delete Nodes And Return Forest

Description

Given the root of a binary tree, each node in the tree has a distinct value.

After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees).

Return the roots of the trees in the remaining forest. You may return the result in any order.

Example 1:

Input: root = [1,2,3,4,5,6,7], to_delete = [3,5]
Output: [[1,2,null,4],[6],[7]]

Example 2:

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

Constraints:

The number of nodes in the given tree is at most 1000.
Each node has a distinct value between 1 and 1000.
to_delete.length <= 1000
to_delete contains distinct values between 1 and 1000.

Solution

If a node is in to_delete, then mark its child node as new root. Also, if a node’s child is in to_delete, delete them.

Use post-order traversal.

Time complexity: $o(n)$
Space complexity: $o(n)$

Code

# 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 delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
        stack = [(root, 0)]
        res = []
        while stack:
            node, status = stack.pop()
            if status == 0:
                stack.append((node, 1))
                if node.left:
                    stack.append((node.left, 0))
                if node.right:
                    stack.append((node.right, 0))
            else:
                if node.left and node.left.val in to_delete:
                    node.left = None
                if node.right and node.right.val in to_delete:
                    node.right = None
                if node.val in to_delete:
                    if node.left:
                        res.append(node.left)
                    if node.right:
                        res.append(node.right)
        if root.val not in to_delete:
            res.append(root)
        return res