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.

• Time : O(n)
• Space : O(1)

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode() : val(0), left(nullptr), right(nullptr) {}
 *     TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
 *     TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
 * };
 */
class Solution {
    unordered_set<int> del;
    vector<TreeNode*> res;
    void solution(TreeNode* node, bool isRoot) {
        if(!node) return;
        if(isRoot) {
            if(!del.count(node->val)) {
                res.push_back(node);
                if(node->left and del.count(node->left->val)) {
                    solution(node->left->left, true);
                    solution(node->left->right, true);
                    node->left = NULL;
                } else solution(node->left, false);

                if(node->right and del.count(node->right->val)) {
                    solution(node->right->left, true);
                    solution(node->right->right, true);
                    node->right = NULL;
                } else solution(node->right, false);
            } else {
                solution(node->left, true);
                solution(node->right, true);
            }
        } else {
            if(node->left and del.count(node->left->val)) {
                solution(node->left->left, true);
                solution(node->left->right, true);
                node->left = NULL;
            } else solution(node->left, false);

            if(node->right and del.count(node->right->val)) {
                solution(node->right->left, true);
                solution(node->right->right, true);
                node->right = NULL;
            } else solution(node->right, false);
        }
    }
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        for(auto& d : to_delete)
            del.insert(d);
        solution(root, true);
        return res;
    }
};