leetcode 1110 - Delete Nodes And Return Forest

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

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
class Solution {
public:
unordered_set<int> set_to_delete;
vector<TreeNode *> result;

vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
for (int i : to_delete)
set_to_delete.insert(i);

delNode(root, true);
return result;
}

TreeNode *delNode(TreeNode *root, bool parent_deleted) {
if (!root) return NULL;

if (set_to_delete.count(root->val)) {
delNode(root->left, true);
delNode(root->right, true);
return NULL;
} else {
root->left = delNode(root->left, false);
root->right = delNode(root->right, false);
if (parent_deleted)
result.push_back(root);
return root;
}
}
};