Problem of The Day: Delete Nodes And Return Forest

Problem Statement

Intuition

I believe the problem can be approached using a depth-first search (DFS) traversal. By recursively visiting each node, I can determine whether it should be deleted and manage the connections to its children accordingly. If a node is marked for deletion, its children (if they are not also marked for deletion) should become new roots in the resulting forest.

Approach

1. Initial Setup: I’ll start by checking if the root itself should be deleted. If not, I’ll add it to the result list.
2. Depth-First Search (DFS): Using DFS, I’ll traverse the tree and handle the deletion process:
   • If a node is in the to_delete list, I’ll check its children. Any child not in to_delete will be added to the result list.
   • I’ll also break the link between the parent and the node marked for deletion.
3. Return the Result: After the DFS traversal, the result list will contain all the new roots of the remaining trees.

Complexity

• Time Complexity: (O(n)), where (n) is the number of nodes in the tree. Each node is visited once.
• Space Complexity: (O(n)), due to the recursion stack in the worst case and the space needed for the result list.

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]:
        res = [root] if root and root.val not in to_delete else []

        def dfs(node, parent):
            if not node:
                return
            dfs(node.left, node)
            dfs(node.right, node)
            if node.val in to_delete:
                if node.left and node.left.val not in to_delete:
                    res.append(node.left)
                if node.right and node.right.val not in to_delete:
                    res.append(node.right)
                if parent and parent.left is node:
                    parent.left = None
                if parent and parent.right is node:
                    parent.right = None

        dfs(root, None)
        return res

Editorial

Approach 1: Recursion (Postorder Traversal)

class Solution:
    def delNodes(
        self, root: Optional[TreeNode], to_delete: List[int]
    ) -> List[TreeNode]:
        to_delete_set = set(to_delete)
        forest = []

        root = self._process_node(root, to_delete_set, forest)

        # If the root is not deleted, add it to the forest
        if root:
            forest.append(root)

        return forest

    def _process_node(
        self, node: TreeNode, to_delete_set: Set[int], forest: List[TreeNode]
    ) -> TreeNode:
        if not node:
            return None

        node.left = self._process_node(node.left, to_delete_set, forest)
        node.right = self._process_node(node.right, to_delete_set, forest)

        # Node Evaluation: Check if the current node needs to be deleted
        if node.val in to_delete_set:
            # If the node has left or right children, add them to the forest
            if node.left:
                forest.append(node.left)
            if node.right:
                forest.append(node.right)
            # Delete the current node by returning None to its parent
            return None

        return node

Approach 2: BFS Forest Formation

class Solution:
    def delNodes(
        self, root: Optional[TreeNode], to_delete: List[int]
    ) -> List[TreeNode]:
        if not root:
            return []

        to_delete_set = set(to_delete)
        forest = []

        nodes_queue = deque([root])

        while nodes_queue:
            current_node = nodes_queue.popleft()

            if current_node.left:
                nodes_queue.append(current_node.left)
                # Disconnect the left child if it needs to be deleted
                if current_node.left.val in to_delete_set:
                    current_node.left = None

            if current_node.right:
                nodes_queue.append(current_node.right)
                # Disconnect the right child if it needs to be deleted
                if current_node.right.val in to_delete_set:
                    current_node.right = None

            # If the current node needs to be deleted, add its non-null children to the forest
            if current_node.val in to_delete_set:
                if current_node.left:
                    forest.append(current_node.left)
                if current_node.right:
                    forest.append(current_node.right)

        # Ensure the root is added to the forest if it is not to be deleted
        if root.val not in to_delete_set:
            forest.append(root)

        return forest