Delete Nodes And Return Forest

Problem Id: 1110 Difficulty: Medium Tag: Tree Tag: Depth-first Search


Intuition

Do a DFS to remove all the to_delete nodes. When a node is removed, its left sub-tree and right sub-tree should be added to the forest.

Solution


# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:
        self.forest = []
        self.to_delete = set(to_delete)
        self._dfs_remove(root, None)
        return self.forest

    def _dfs_remove(self, node, parent):
        if node is None:
            return None
        if node.val in self.to_delete:
            self._dfs_remove(node.left, None)
            self._dfs_remove(node.right, None)
            return None
        elif parent is None:
            self.forest.append(node)
        node.left = self._dfs_remove(node.left, node)
        node.right = self._dfs_remove(node.right, node)
        return node