Delete Nodes And Return Forest

medium tree dfs recursion

Problem

Given the root of a binary tree with unique values and an array to_delete, delete those nodes and return the roots of the resulting forest.

Inputroot = [1,2,3,4,5,6,7], to_delete = [3,5]

Output[[1,2,4],[6],[7]]

Deleting 3 promotes 6 and 7 to roots; deleting 5 disconnects no remaining tree.

tree (level-order, 'null' for missing)

to_delete

def del_nodes(root, to_delete):
    D = set(to_delete)
    forest = []
    def dfs(node, is_root):
        if not node: return None
        deleted = node.val in D
        if is_root and not deleted:
            forest.append(node)
        node.left = dfs(node.left, deleted)
        node.right = dfs(node.right, deleted)
        return None if deleted else node
    dfs(root, True)
    return forest

function delNodes(root, toDelete) {
  const D = new Set(toDelete);
  const forest = [];
  function dfs(node, isRoot) {
    if (!node) return null;
    const del = D.has(node.val);
    if (isRoot && !del) forest.push(node);
    node.left = dfs(node.left, del);
    node.right = dfs(node.right, del);
    return del ? null : node;
  }
  dfs(root, true);
  return forest;
}

class Solution {
    Set<Integer> D;
    List<TreeNode> forest;
    public List<TreeNode> delNodes(TreeNode root, int[] toDelete) {
        D = new HashSet<>();
        for (int x : toDelete) D.add(x);
        forest = new ArrayList<>();
        dfs(root, true);
        return forest;
    }
    TreeNode dfs(TreeNode node, boolean isRoot) {
        if (node == null) return null;
        boolean del = D.contains(node.val);
        if (isRoot && !del) forest.add(node);
        node.left = dfs(node.left, del);
        node.right = dfs(node.right, del);
        return del ? null : node;
    }
}

vector<TreeNode*> delNodes(TreeNode* root, vector<int>& toDelete) {
    unordered_set<int> D(toDelete.begin(), toDelete.end());
    vector<TreeNode*> forest;
    function<TreeNode*(TreeNode*, bool)> dfs = [&](TreeNode* node, bool isRoot) -> TreeNode* {
        if (!node) return nullptr;
        bool del = D.count(node->val);
        if (isRoot && !del) forest.push_back(node);
        node->left = dfs(node->left, del);
        node->right = dfs(node->right, del);
        return del ? nullptr : node;
    };
    dfs(root, true);
    return forest;
}

Explanation

When you delete some nodes from a tree, the tree breaks into several smaller trees — a forest. The goal is to return the root of each surviving piece. The neat observation is that a node becomes a new root exactly when its parent was deleted (or it is the original root).

So we carry one extra flag through the recursion: isRoot, meaning "could this node be the top of its own tree right now?". We also put the values to remove into a set D for instant lookups.

In dfs(node, isRoot) we first check whether this node is being deleted. If it survives and isRoot is true, we add it to forest. Then we recurse into both children, passing down del as their isRoot flag — because if this node is deleted, its children are about to be promoted to roots.

Finally we return null for a deleted node (so the parent's link drops it) or the node itself if it stays.

Example: root = [1,2,3,4,5,6,7] with to_delete = [3,5]. Deleting 3 frees its children 6 and 7 as new roots, deleting 5 removes a leaf, and the answer is the three trees [[1,2,4],[6],[7]].

Time: O(n) Space: O(n)