1110 Delete Nodes And Return Forest

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

solution

class Solution:
    def __init__(self):
        self.res = []

    def delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
        if self.dfs(root, to_delete):
            self.res.append(root)
        return self.res

    def dfs(self, root, to_delete):
        if not root:
            return

        root.left = self.dfs(root.left, to_delete)
        root.right = self.dfs(root.right, to_delete)

        if root.val in to_delete:
            # 通过delete节点返回None来控制初始根节点是否需要添加。因为这里有个返回，导致左右dfs需要在其上面 -> 后序遍历
            if root.left:
                self.res.append(root.left)
            if root.right:
                self.res.append(root.right)
            return

        return root

时间复杂度：O(N) 空间复杂度：O(N)