Delete Nodes And Return Forest

LeetCode Q 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]]

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:

Solution 1: Using maps

We can first build two maps, one is parentsMap (key: node (TreeNode); value: its parent (TreeNode)), another is nodesMap (key: node value (int); value: correspongding node (TreeNode)).

Code:

Set<TreeNode> set;
Map<Integer, TreeNode> nodesMap;
Map<TreeNode, TreeNode> parentsMap;

public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
  if (root == null) return new ArrayList();
  
  set = new HashSet<>();
  set.add(root);
  
  nodesMap = new HashMap<>();
  
  parentsMap = new HashMap<>();
  parentsMap.put(root, root);
  
  buildMaps(root);
  
  for (int nodeVal: to_delete) {
    TreeNode node = nodesMap.get(nodeVal);
    
    if (set.contains(node)) set.remove(node);
    
    TreeNode p = parentsMap.get(node);
    
    if (p.left == node) p.left = null;
    if (p.right == node) p.right = null;
    
    if (node.left != null) set.add(node.left);
    
    if (node.right != null) set.add(node.right);
  }
  
  return new ArrayList(set);
}
    
private void buildMaps(TreeNode root) {
  nodesMap.put(root.val, root);

  if (root.left != null) {
    parentsMap.put(root.left, root);
    buildMaps(root.left);
  }

  if (root.right != null) {
    parentsMap.put(root.right, root);
    buildMaps(root.right);
  }
}

Solution 2: bottom-up DFS

Code:

public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
  List<TreeNode> forest = new ArrayList<>();
  if (root == null) return forest;

  Set<Integer> set = new HashSet<>();
  for (int d: to_delete) set.add(d);

  deleteNodes(root, set, forest);

  if (!set.contains(root.val)) forest.add(root);

  return forest;
}

private TreeNode deleteNodes(TreeNode root, Set<Integer> set, List<TreeNode> forest) {
  if (root == null) return null;

  node.left = deleteNodes(node.left, set, forest);
  node.right = deleteNodes(node.right, set, forest);

  if (set.contains(root.val)) {
    if (root.left != null) forest.add(root.left);
    if (root.right != null) forest.add(root.right);
  	return null;
  }

  return root;
}