Problem statement

https://leetcode.com/problems/delete-nodes-and-return-forest/

Solution

We need to use postorder dfs: first we traverse left and right subtrees and return roots of these subtrees. Then we look at node and decide what to do: if it needs to be deleted, we add to our answer roots of left and right subtrees if they exist and return None. If it is not to be deleted return node itself. In the end we also check the root of our tree.

Complexity

It is O(n) for time and O(n) for space.

Code

class Solution:
    def delNodes(self, root, D):
        def dfs(node):
            if not node: return node
            node.left = dfs(node.left)
            node.right = dfs(node.right)
			
            if node.val in D:
                for child in [node.left, node.right]:
                    if child: self.ans += [child]
                return None
            return node
        
        self.ans, D = [], set(D)
        dfs(root)
        if root.val not in D:
            self.ans += [root]
        return self.ans