Solution: Delete Nodes And Return Forest

Let’s solve the Delete Nodes And Return Forest problem using the Tree Depth-first Search pattern.

Statement

Given the root of a binary tree where each node has a unique value, your task is to delete all nodes with values specified in the delete_nodes list. After performing the deletions, the tree will split into a forest—a collection of disjoint trees. Return the roots of the remaining trees in the forest in any order.

Constraints:

  • 00\leq nodes 100\leq100

  • 11\leq nodes.value 1000\leq 1000

  • 00\leq delete_nodes.length 100\leq 100

  • 11\leq delete_nodes[i] 1000\leq 1000

Note: Both nodes and delete_nodes[i] will have distinct values.

Solution

When certain nodes are removed from a binary tree, the remaining parts of the tree form disjoint subtrees, often referred to as a forest. The key insight is that when a node is deleted, its children, if any, become the roots of new trees. The algorithm must ensure that each node is either included in its current tree or, if its parent is deleted, becomes the root of a new tree. This process is handled iteratively as the tree is traversed, ensuring all nodes are processed correctly.

Let’s look at the step-by-step implementation of the solution:

  • Check if the root is None, and return an empty list, as no nodes exist to process.

  • Convert the delete_nodes list into a set for efficient lookup.

  • Initialize an empty list to store the roots of the resulting forest.

  • Use a stack to keep track of nodes to be processed, starting with the root node.

  • While the stack is not empty, pop a node from the stack and perform the following steps:

    • If the left child exists:

      • Push it onto the stack for further processing.

      • Check if the left child’s data is in the to_delete set.

        • If it is, set the left child to None to detach it from the current node.

    • If the right child exists:

      • Push it onto the stack for further processing.

      • Check if the right child’s data is in the to_delete set.

        • If it is, set the right child to None to detach it from the current node.

    • If the current node’s data is in the to_delete set.

      • If the left child exists, add it to the forest list as a new root.

      • If the right child exists, add it to the forest list as a new root.

  • After processing all nodes, check if the original root node’s data is not in the to_delete set.

    • If it isn’t, add the root to the forest list.

  • Return the forest list containing the remaining tree roots after the specified nodes have been deleted.

Let’s look at the following illustration to get a better understanding of the solution:

Level up your interview prep. Join Educative to access 70+ hands-on prep courses.