Leetcode 1110. Delete Nodes And Return Forest (Python)

Related Topic

Depth-First-Search.

Description

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.

Sample I/O

Example 1

Input: root = [1,2,3,4,5,6,7], to_delete = [3,5]
Output: [[1,2,null,4],[6],[7]]

Methodology

Use DFS to traversal the tree in postorder, if the node does not need to delete and the node return the node, if the node need to delete and the node has children, then append children into the res list first and return None.

Code(DFS)

def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:
        ds = set(to_delete)
        res = []
        def dfs(root):
            if root is None: return
            root.left = dfs(root.left)
            root.right = dfs(root.right)
            if root.val not in ds:return root
            if root.left: res.append(root.left)
            if root.right: res.append(root.right)
            return None
        root = dfs(root)
        if root: res.append(root)
        return res

BigO

We traversal all elements of the tree once so total time complexity is O(n)