1110.delete-nodes-and-return-forest.md

1110. Delete Nodes And Return Forest

Clarification Questions

• Does the to_delete array contain the node that does not exist in the tree?

Test Cases

Normal Cases

Input: 
        1               
     /     \      
    3       5          
  /   \    / \
 8    9   7   0   
         /   / 
        2   6   
delete = [5]
Output: 
        1               
     /          
    3                 
  /   \    
 8    9   7   0   
         /   / 
        2   6  

Edge / Corner Cases

• The nodes to delete are in the same subtree.

Input: 
        1               
     /     \      
    3       5          
  /   \    / \
 8    9   7   0   
         /   / 
        2   6  
delete = [1,5]
Output: 
    3               
  /   \    
 8    9   7   0   
         /   / 
        2   6  

• The root is deleted.
• All the nodes are deleted.
• All the nodes in subtrees are deleted.

Breakdown

Given the values to delete, how to delete the nodes (prune all the subtrees) and return the root?

fun del(root: TreeNode?, toDelete: Set<Int>): TreeNode? {
	if (root == null) return null
	if (root.`val` in toDelete) {
		return null
	}
	root.left = del(root.left, toDelete)
	root.right = del(root.right, toDelete)
	return root
}

Postorder

Idea!! We traverse each node as root, and check if we should delete the root. If we should delete the root, we add the children to the forest. But bfore that, we need to check all the children first.

      root       del(root)
     /    \
  left   right
  ...       ...

// After delete `root`, the subtrees `left` and `right` become the forest.
     /    \
  left   right
  ...       ...

We use postorder to delete, because we need to delete the nodes from the bottom to the top.

Here is one key point to note, if the original root is not in the delete list, we should add it to the forest first.

fun delNodes(root: TreeNode?, to_delete: IntArray): List<TreeNode?> {
    val set = to_delete.toHashSet()
    val forest = mutableListOf<TreeNode?>()
    val updatedRoot = delete(root, set, forest)
    // If the root is not deleted, we should add it to the forest. Otherwise, the root won't be added to the forest.
    if (updatedRoot != null) forest.add(updatedRoot)
    return forest
}

private fun delete(root: TreeNode?, deleteSet: HashSet<Int>, forest: MutableList<TreeNode?>): TreeNode? {
    if (root == null) return null

	// We delete the child nodes first.
    root.left = delete(root.left, deleteSet, forest)
    root.right = delete(root.right, deleteSet, forest)

	// Then check if we should delete the current root, then all children become the forest.
    if (deleteSet.contains(root.`val`)) {
        if (root.left != null) forest.add(root.left)
        if (root.right != null) forest.add(root.right)
        return null
    }
    return root
}

• Time Complexity: O(n)
• Space Complexity: O(n)