Delete Nodes And Return Forest Python Solution

lcpy gen -n 1110   # by problem number
lcpy gen -s delete_nodes_and_return_forest   # by problem name

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

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

from leetcode_py import TreeNode


class Solution:
    # Time: O(n)
    # Space: O(n)
    def del_nodes(self, root: TreeNode[int] | None, to_delete: list[int]) -> list[TreeNode[int]]:
        delete = set(to_delete)
        forest: list[TreeNode[int]] = []

        def dfs(node: TreeNode[int] | None, is_root: bool) -> TreeNode[int] | None:
            if node is None:
                return None
            deleted = node.val in delete
            if is_root and not deleted:
                forest.append(node)
            node.left = dfs(node.left, deleted)
            node.right = dfs(node.right, deleted)
            return None if deleted else node

        dfs(root, True)
        return forest