1110. Delete Nodes And Return Forest #

Problem #

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.

Problem Summary #

Given the root node root of a binary tree, each node in the tree has a distinct value. If a node’s value appears in to_delete, we delete that node from the tree, and finally obtain a forest (a collection of some disjoint trees). Return each tree in the forest. You may organize the answer in any order.

Notes:

• The maximum number of nodes in the tree is 1000.
• Each node has a value between 1 and 1000, and all values are distinct.
• to_delete.length <= 1000
• to_delete contains some distinct values from 1 to 1000.

Solution Approach #

• Given a tree and an array, delete the nodes whose values are the same as the elements in the array. Output the forest after all deletions.
• Simple problem. Traverse the tree while deleting the elements in the array. Here, we can first put the elements in the array into a map to speed up lookup. Delete the node when encountering the same element. The special case that needs to be handled here is whether the currently deleted node is the root node; if it is the root node, its left child or right child needs to be set to nil according to the condition.

Code #

func delNodes(root *TreeNode, toDelete []int) []*TreeNode {
	if root == nil {
		return nil
	}
	res, deleteMap := []*TreeNode{}, map[int]bool{}
	for _, v := range toDelete {
		deleteMap[v] = true
	}
	dfsDelNodes(root, deleteMap, true, &res)
	return res
}

func dfsDelNodes(root *TreeNode, toDel map[int]bool, isRoot bool, res *[]*TreeNode) bool {
	if root == nil {
		return false
	}
	if isRoot && !toDel[root.Val] {
		*res = append(*res, root)
	}
	isRoot = false
	if toDel[root.Val] {
		isRoot = true
	}
	if dfsDelNodes(root.Left, toDel, isRoot, res) {
		root.Left = nil
	}
	if dfsDelNodes(root.Right, toDel, isRoot, res) {
		root.Right = nil
	}
	return isRoot
}