Leetcode 1207 Solution

This article provides solution to leetcode question 1207 (delete-nodes-and-return-forest)
Categories:
Tags:
Access this page by simply typing in "lcs 1207" in your browser address bar if you have bunnylol configured.

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

Solution

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:
        ans = []

        to_delete = set(to_delete)

        def del_nodes(node, to_add):
            if node is None:
                return

            if node.val in to_delete:
                node.left = del_nodes(node.left, True)
                node.right = del_nodes(node.right, True)
                return None
            else:
                if to_add:
                    ans.append(node)

                node.left = del_nodes(node.left, False)
                node.right = del_nodes(node.right, False)
                return node

        del_nodes(root, True)

        return ans