1110. Delete Nodes And Return Forest

/**

 * 1110. Delete Nodes And Return Forest 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]]

 */

See the solution at github here: https://github.com/zcoderz/leetcode/blob/main/src/main/java/google/medium/DeleteNodesAndReturnForest.java

See another similar question that requests to delete nodes in a BST also included in solutions.

Observation: We need to delete items from a tree and add subtrees formed into a result array. We will need to traverse the tree and start this process from bottom up so that as we process parent nodes, we are guaranteed by design that the child nodes have already been processed.

Logic:

  1. Traverse the BST recursively
  2. When returning from recursion call stack
    1. Check for left child if it was in the set of values to be deleted
      1. If it was then add its children into sub trees array
    1. Repeat for right child

Here is the code:

public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
    Set<Integer> set = new HashSet<>();
    for (int v : to_delete) {
        set.add(v);
    }

    //to do figure out how to handle delete of root
    traverseTree(root, set);
    if (set.contains(root.val)) {
        if (root.left != null) {
            roots.add(root.left);
        }
        if (root.right != null) {
            roots.add(root.right);
        }
    } else {
        roots.add(root);
    }
    return roots;
}

void traverseTree(TreeNode node, Set<Integer> toDelete) {
    if (node == null) {
        return;
    }

    traverseTree(node.left, toDelete);
    traverseTree(node.right, toDelete);

    if ((node.left != null) && (toDelete.contains(node.left.val))) {
        toDelete.remove(node.left.val);
        if (node.left.left != null) {
            roots.add(node.left.left);
        }
        if (node.left.right != null) {
            roots.add(node.left.right);
        }
        node.left = null;
    }

    if ((node.right != null) && (toDelete.contains(node.right.val))) {
        toDelete.remove(node.right.val);
        if (node.right.left != null) {
            roots.add(node.right.left);
        }
        if (node.right.right != null) {
            roots.add(node.right.right);
        }
        node.right = null;
    }
}