Delete Nodes And Return Forest LeetCode solution in C++

1
class Solution {
2
  void solve(TreeNode *root, TreeNode *prev, unordered_map<int, int> &mp, int left,
3
             vector<TreeNode *> &ans) {
4
    if (root == NULL) return;
5

6
    int f = 0;
7

8
    solve(root->left, root, mp, 1, ans);
9
    solve(root->right, root, mp, 0, ans);
10

11
    if (mp.find(root->val) != mp.end()) {
12
      f = 1;
13
      if (root->right) ans.push_back(root->right);
14
      // cout<<root->left->val<<" ";}
15
      if (root->left) ans.push_back(root->left);
16
      // cout<<root->right->val<<" ";}
17
    }
18

19
    if (f == 1) {
20
      if (left)
21
        prev->left = NULL;
22
      else
23
        prev->right = NULL;
24
    }
25
  }
26

27
 public:
28
  vector<TreeNode *> delNodes(TreeNode *root, vector<int> &to_delete) {
29
    if (root == NULL) return {};
30

31
    unordered_map<int, int> mp;
32
    for (int i = 0; i < to_delete.size(); i++) mp[to_delete[i]]++;
33

34
    vector<TreeNode *> ans;
35

36
    TreeNode *dummy = new TreeNode(-1);
37
    dummy->left = root;
38

39
    solve(dummy->left, dummy, mp, 1, ans);
40
    if (dummy->left) ans.push_back(root);
41

42
    return ans;
43
  }
44
};
45
```