Delete Nodes And Return Forest LeetCode solution in Python

1
# Definition for a binary tree node.
2
# class TreeNode:
3
#     def __init__(self, val=0, left=None, right=None):
4
#         self.val = val
5
#         self.left = left
6
#         self.right = right
7
class Solution:
8
    def traverse(self, node, par):
9
        if node:
10
            self.parent[node.val] = par
11
            self.traverse(node.left, node)
12
            self.traverse(node.right, node)
13

14
    def deleteNode(self, toDelete):
15
        node = None
16
        par = self.parent[toDelete]
17
        if par.val == toDelete:
18
            node = par
19
        elif par.left and par.left.val == toDelete:
20
            node = par.left
21
        elif par.right and par.right.val == toDelete:
22
            node = par.right
23
        if node.left:
24
            self.unique[node.left] = True
25
        if node.right:
26
            self.unique[node.right] = True
27

28
        if node in self.unique:
29
            self.unique.pop(node)
30
        if node != self.parent[toDelete]:
31
            if par.left == node:
32
                par.left = None
33
            else:
34
                par.right = None
35

36
    def delNodes(
37
        self, root: Optional[TreeNode], to_delete: List[int]
38
    ) -> List[TreeNode]:
39
        self.parent = {}
40
        self.traverse(root, root)
41
        self.unique = {root: True}
42
        for node in to_delete:
43
            self.deleteNode(node)
44
        return self.unique