LeetCode 1110. Delete Nodes And Return Forest

Description

https://leetcode.com/problems/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.

Example 1:

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

Example 2:

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

Constraints:

• The number of nodes in the given tree is at most 1000.
• Each node has a distinct value between 1 and 1000.
• to_delete.length <= 1000
• to_delete contains distinct values between 1 and 1000.

Python Solution

Traverse the tree, if the node value is in the delete value list, change the node to be null, also check if the node has sub tree or not if there is a subtree, then mark it as a separate disjoint tree.

# 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]:
        results = []

        if not root:
            return results
                        
        to_delete = set(to_delete)
        
        self.helper(root, to_delete, results, True)
        
        return results
        
        
    def helper(self, root, to_delete, results, is_root):
        if not root:
            return None
        
        if root.val in to_delete:
            if root.left:
                left = self.helper(root.left, to_delete, results, True)
            if root.right:
                right = self.helper(root.right, to_delete, results, True)

            root = None
            
            return root
        else:
            left = self.helper(root.left, to_delete, results, False)
            right = self.helper(root.right, to_delete, results, False)


            root.left = left
            root.right = right
            
            if is_root:
                results.append(root)
            
            return root