Taro Logo

Delete Nodes And Return Forest

Medium
Meta logo
Meta
Topics:
TreesRecursion

You are given the root of a binary tree, where each node in the tree has a distinct value.

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

Can you return the roots of the trees in the remaining forest? You may return the result in any order.

Example 1:

Consider the following binary tree:

    1
   / \
  2   3
 / \
4   5

If to_delete = [3, 5], the expected output is [[1, 2, null, 4], [6], [7]].

Example 2:

Consider the following binary tree:

    1
   / \
  2   4
   \  
    3

If to_delete = [3], the expected output is [[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


Let's analyze the problem. We're given a binary tree and a list of node values to delete. After deleting these nodes, we're left with a forest, which is a collection of disjoint trees. The goal is to return the roots of these remaining trees.

Naive Approach

A brute-force approach would involve the following steps:

  1. Traverse the tree and mark nodes to be deleted.
  2. Perform another traversal to identify the roots of the remaining trees. A node is a root if its parent was deleted or if it's the original root and wasn't deleted.

This approach involves multiple traversals of the tree, resulting in potentially higher time complexity.

Optimal Approach

A more efficient approach involves a single traversal of the tree using recursion. During the traversal, we check if the current node's value is in the to_delete list. If it is, we remove the node and add its children (if they exist) to the result list as new roots. If it isn't, we continue the traversal, ensuring that the parent-child relationship is maintained correctly.

Algorithm:

  1. Create a HashSet from the to_delete array for efficient lookups.
  2. Create a list to store the roots of the remaining forest.
  3. Define a recursive helper function that takes a node and the HashSet as input.
  4. In the helper function:
    • Base case: If the node is null, return null.
    • Recursively process the left and right subtrees, updating the left and right children of the current node.
    • Check if the current node's value is in the HashSet.
      • If it is, add the left and right children (if they exist) to the result list. Return null to remove the current node from the tree.
      • If it isn't, return the current node.
  5. Call the helper function with the root of the tree. If the root is not deleted, add it to the result list.
  6. Return the result list.

Edge Cases:

  • Empty tree: If the root is null, return an empty list.
  • All nodes to be deleted: If all nodes are in to_delete, return an empty list.
  • No nodes to be deleted: If to_delete is empty, return a list containing the original root.
  • Root needs to be deleted: Ensure the root is checked and its children (if any) are added to the result.

Code (Java):

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

class TreeNode {
    int val;
    TreeNode left;
    TreeNode right;

    TreeNode() {}

    TreeNode(int val) {
        this.val = val;
    }

    TreeNode(int val, TreeNode left, TreeNode right) {
        this.val = val;
        this.left = left;
        this.right = right;
    }
}

class Solution {
    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        List<TreeNode> result = new ArrayList<>();
        Set<Integer> toDeleteSet = new HashSet<>();
        for (int val : to_delete) {
            toDeleteSet.add(val);
        }

        root = helper(root, toDeleteSet, result);
        if (root != null) {
            result.add(root);
        }

        return result;
    }

    private TreeNode helper(TreeNode node, Set<Integer> toDeleteSet, List<TreeNode> result) {
        if (node == null) {
            return null;
        }

        node.left = helper(node.left, toDeleteSet, result);
        node.right = helper(node.right, toDeleteSet, result);

        if (toDeleteSet.contains(node.val)) {
            if (node.left != null) {
                result.add(node.left);
            }
            if (node.right != null) {
                result.add(node.right);
            }
            return null;
        } else {
            return node;
        }
    }
}

Time Complexity:

  • O(N), where N is the number of nodes in the tree. This is because we visit each node once during the traversal.

Space Complexity:

  • O(H + D), where H is the height of the tree (for the recursion stack) and D is the number of elements in to_delete (for the HashSet).
  • In the worst case, if the tree is skewed, H can be O(N).