class Solution {
public:
TreeNode* process(TreeNode* root, unordered_set<int>& set, vector<TreeNode*>& res)
{
if (!root) return NULL;
root->left = process(root->left, set, res);
root->right = process(root->right, set, res);
if (set.find(root->val) != set.end()) {
if (root->left){
res.push_back(root->left);
}
if (root->right){
res.push_back(root->right);
}
delete(root);
return NULL;
}
return root;
}
vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
vector<TreeNode*> res;
unordered_set<int> set(to_delete.begin(), to_delete.end());
if(process(root, set, res)!=NULL)
{
res.push_back(root);
}
return res;
}
};
|