LeetCode #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.

class Solution {
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        for(int val:to_delete) s.insert(val);
        vector<TreeNode*> result;
        traverse(root,NULL,result);
        return result;
    }
    
    void traverse(TreeNode* root, TreeNode* pre, vector<TreeNode*>& result)
    {
        // 遍历二叉树，前一个节点被删除，同时当前节点不需要被删除时，加入结果中
        if(root==NULL) return;
        if(pre==NULL) // 先判断根节点，排除pre为空的情况
        {
            if(s.count(root->val)==0)
                result.push_back(root);
        }
        // pre不为空，而且当前节点需要删除
        else if(s.count(root->val))
        {
            if(pre->left==root) pre->left=NULL;
            else pre->right=NULL;   
        }
        // pre不为空，而且当前节点不需要删除，加入结果中
        else if(s.count(pre->val)) result.push_back(root); 
        traverse(root->left,root,result);
        traverse(root->right,root,result);   
    }
    
private:
    unordered_set<int> s;
};