leetcode 1110.删点成林

 

/**
 * 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) {}
 * };
 */
vector<TreeNode*> ans;
int n;
bool flag[1007];

bool dfs(TreeNode* now){
    if (now==NULL)
        return false;
    if (dfs(now->left)) //下层被删去
        now->left=NULL; 
    if (dfs(now->right))
        now->right=NULL;
    if (flag[now->val]){
        if ((now->left)!=NULL)
            ans.push_back(now->left);
        if ((now->right)!=NULL)
            ans.push_back(now->right);
        return true;
    }
    return false;
}

class Solution {
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        ans.clear();
        n=to_delete.size();
        memset(flag,0,sizeof(flag));
        bool check=false;
        for (int i=0;i<n;i++){
            flag[to_delete[i]]=true;
            if (to_delete[i]==(root->val))
                check=true;
        }
        if (!check)
            ans.push_back(root);
        dfs(root);   
        return ans;
    }
};