aboutsummaryrefslogtreecommitdiff
path: root/delete_nodes_return_forest.py
blob: d9a95ad0faabe64fced241efecb2fefa41868913 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
# 1110. Delete Nodes And Return Forest

"""
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.
"""


# Definition for a binary tree node.
class TreeNode(object):
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution(object):
    def delNodes(self, root, to_delete):
        """
        :type root: TreeNode
        :type to_delete: List[int]
        :rtype: List[TreeNode]
        """
        sel_del, ans = set(to_delete), []

        def helper(root, is_root):
            if not root:
                return None
            root_del = root.val in sel_del
            if is_root and not root_del:
                ans.append(root)
            root.left = helper(root.left, root_del)
            root.right = helper(root.right, root_del)
            return None if root_del else root

        helper(root, True)
        return ans