Leetcode/src/1110-Delete-Nodes-And-Return-Forest/1110.cpp at master · fly-chosen/Leetcode · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
class Solution
{
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete)
    {
        delete_set = unordered_set<int>(to_delete.begin(), to_delete.end());
        preOrder(root, true);
        return res;
    }
private:
    unordered_set<int> delete_set;
    vector<TreeNode*> res;

    TreeNode* preOrder(TreeNode* root, bool n_root)
    {
        if (root == nullptr) return nullptr;
        bool root_delete = delete_set.count(root->val) > 0;
        if (!root_delete and n_root) res.push_back(root);
        root->left = preOrder(root->left, root_delete);
        root->right = preOrder(root->right, root_delete);
        return root_delete ? nullptr : root;
    }
};