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.

 

Example 1:

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

 

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.

思路：递归，左右子树收集信息，感觉自己写的有点复杂

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


class Solution(object):
    def delNodes(self, root, to_delete):
        """
        :type root: TreeNode
        :type to_delete: List[int]
        :rtype: List[TreeNode]
        """
        def helper(root):
            if not root: return []
            ls = helper(root.left)
            rs = helper(root.right)
            
            rl,rr=root.left,root.right
            root.left,root.right=None,None
            
            if root.val not in to_delete:
                merge=[]
                for l in ls:
                    if l.val!=rl.val:
                        merge.append(l)
                    else:
                        root.left=l
                for r in rs:
                    if r.val!=rr.val:
                        merge.append(r)
                    else:
                        root.right=r
                merge.append(root)
            else:
                merge = ls+rs
                
            return merge
        
        return helper(root)