Leetcode-1110. Delete Nodes and Return Forest

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

First, consider a similar but replacing trees to linked node question,

There are two kinds of nodes:
To be deleted (D)
and normal (N)

There are 4 relations between node and node.next (left/right):
D -> D
D -> N
N -> D
N -> N

For DtoD and NtoN, we don’t need to be worried. We can simply move the pointer to the next.

For DtoN, we need to do

1. delete D.next
2. append N to the list

For NtoD, we need to do

1. delete N.next

Replacing all above “next” to left and right, we already have a solution.

What’s more, for finding convinence, I cast to_delete into a set.

Code:

class Solution(object):
    def delNodes(self, root, to_delete):
        """
        :type root: TreeNode
        :type to_delete: List[int]
        :rtype: List[TreeNode]
        """
        
        global s 
        s = set(to_delete)
        
        padNode = TreeNode(float("Inf"))
        padNode.left = root
        s.add(float("Inf"))
        
        root = padNode
 
        op = []
        
        op.extend(self.helper(root))
        
        return op
    
    def helper(self, root):

        op = []
        
        if root.val in s:
            
            if root.left:
                if root.left.val in s:
                    op.extend(self.helper(root.left))
                else:
                    op.append(root.left)
                    op.extend(self.helper(root.left))
                    root.left = None
            if root.right:
                if root.right.val in s:
                    op.extend(self.helper(root.right))
                else:
                    op.append(root.right)
                    op.extend(self.helper(root.right))
                    root.right = None
        else:
            
            if root.left:
            
                op.extend(self.helper(root.left))
                if root.left.val in s:
                    root.left = None
            if root.right:
                op.extend(self.helper(root.right))
                if root.right.val in s:
                    root.right = None
                    
        return op