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.

TreeNode* helper(TreeNode* node, unordered_set<int>& to_delete_st, vector<TreeNode*>& res, bool is_root) {
    if (!node) return nullptr;
    bool deleted = to_delete_st.count(node->val);
    if (is_root && !deleted) res.push_back(node);
    // If the current node is deleted, its children will become new roots for other subtrees
    node->left = helper(node->left, to_delete_st, res, deleted);
    node->right = helper(node->right, to_delete_st, res, deleted);
    return deleted ? nullptr : node;
}
vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
    unordered_set<int> to_delete_st(to_delete.begin(), to_delete.end());
    vector<TreeNode*> res;
    helper(root, to_delete_st, res, true);
    return res;
}

TreeNode* helper(TreeNode* node, unordered_set<int>& to_delete_st, vector<TreeNode*>& res) {
    if (!node) return nullptr;
    node->left = helper(node->left, to_delete_st, res);
    node->right = helper(node->right, to_delete_st, res);
    if (to_delete_st.count(node->val)) {
        if (node->left) {
            res.push_back(node->left);
            node->left = nullptr;
        }
        if (node->right) {
            res.push_back(node->right);
            node->right = nullptr;
        }
        return nullptr;
    } else {
        return node;
    }
}
vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
    unordered_set<int> to_delete_st(to_delete.begin(), to_delete.end());
    vector<TreeNode*> res;
    if (!to_delete_st.count(root->val)) res.push_back(root);
    helper(root, to_delete_st, res);
    return res;
}