删点成林

这道题出自LeetCode，题目如下：

给出二叉树的根节点 root，树上每个节点都有一个不同的值。

如果节点值在 to_delete 中出现，我们就把该节点从树上删去，最后得到一个森林（一些不相交的树构成的集合）。

返回森林中的每棵树。你可以按任意顺序组织答案。

示例：

输入：root = [1,2,3,4,5,6,7], to_delete = [3,5]
输出：[[1,2,null,4],[6],[7]]

提示：

• 树中的节点数最大为 1000。
• 每个节点都有一个介于 1 到 1000 之间的值，且各不相同。
• to_delete.length <= 1000
• to_delete 包含一些从 1 到 1000、各不相同的值。

这道题显然是跟树遍历有关的问题，因为这里需要考虑到节点之间的父子关系，所以采用DFS遍历比较合适。针对这道题，DFS又分前序遍历（即先遍历根节点）和后序遍历（即最后遍历根节点）两种方式，让我们先从前序遍历看起。首先，根据题目假设，每个节点的值是不重复的，所以我们可以将要删除的节点存到set中去，这样可以加快查询速度：

set<int> keys;

for(int k : to_delete)
{
    keys.insert(k);
}

其次，问题的重点在于先序遍历如何定位到可以被加入到森林中的节点。通过观察不难发现，如果当前遍历到的节点不在删除列表中，且它的父节点在删除列表中，它才可以被加入到森林中去：

void dfs(TreeNode* root, set<int>& keys, bool pDelete, vector<TreeNode*> &res)
{
    if(!root)
    {
        return;
    }

    bool rDelete = keys.count(root->val) != 0;
    if(pDelete && !rDelete)
    {
        res.push_back(root);
    }

    dfs(root->left, keys, rDelete, res);
    dfs(root->right, keys, rDelete, res);
}

代码中，pDelete代表当前节点的父节点是否要被删除，rDelete代表当前节点是否要被删除。

除此之外，我们还需要解决一个问题，就是要修改二叉树的结构。当一个节点被删除时，它的父节点对应指向它的left指针或者right指针要置为空。这里有个比较巧妙的做法，就是让递归函数返回一个指针，表示修改过后的节点，如果这个节点被删除了，那么返回的就是空指针，这样上层调用该函数时，直接覆盖对应的left指针或者right指针即可：

TreeNode* dfs(TreeNode* root, set<int>& keys, bool pDelete, vector<TreeNode*> &res)
{
    if(!root)
    {
        return nullptr;
    }

    bool rDelete = keys.count(root->val) != 0;
    if(pDelete && !rDelete)
    {
        res.push_back(root);
    }

    root->left = dfs(root->left, keys, rDelete, res);
    root->right = dfs(root->right, keys, rDelete, res);

    if(rDelete)
    {
        return nullptr;
    }
    else
    {
        return root;
    }
}

返回指针还有一个好处，就是在第一次二叉树的根节点调用该函数时，得到的是修改过的根节点指针，如果指针为空，说明根节点被删除了，就无需将根节点加入到最后的森林列表中：

root = dfs(root, keys, false, res);
if(root)
{
	res.push_back(root);
}

最终通过的代码如下：

/**
 * 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 {
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        vector<TreeNode*> res;

        set<int> keys;

        for(int k : to_delete)
        {
            keys.insert(k);
        }

        root = dfs(root, keys, false, res);
        if(root)
        {
            res.push_back(root);
        }
        return res;
    }

    TreeNode* dfs(TreeNode* root, set<int>& keys, bool pDelete, vector<TreeNode*> &res)
    {
        if(!root)
        {
            return nullptr;
        }

        bool rDelete = keys.count(root->val) != 0;
        if(pDelete && !rDelete)
        {
            res.push_back(root);
        }

        root->left = dfs(root->left, keys, rDelete, res);
        root->right = dfs(root->right, keys, rDelete, res);

        if(rDelete)
        {
            return nullptr;
        }
        else
        {
            return root;
        }
    }
};

后序遍历的方式其实就更简单了，如果当前的节点是需要删除的，那么就把它的子节点加入到森林中。因为后序遍历先遍历的是它的子节点，所以如果子节点也被删除了，那么返回回来的就是空指针，就不必加入到森林里。

最终通过的代码如下：

/**
 * 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 {
public:
    vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
        vector<TreeNode*> res;

        set<int> keys;

        for(int k : to_delete)
        {
            keys.insert(k);
        }

        root = dfs(root, keys, res);
        if(root)
        {
            res.push_back(root);
        }
        return res;
    }

    TreeNode* dfs(TreeNode* root, set<int>& keys, vector<TreeNode*> &res)
    {
        if(!root)
        {
            return nullptr;
        }

        root->left = dfs(root->left, keys, res);
        root->right = dfs(root->right, keys, res);

        bool rDelete = keys.count(root->val) != 0;
        if(rDelete)
        {
            if(root->left)
            {
                res.push_back(root->left);
            }
            if(root->right)
            {
                res.push_back(root->right);
            }
            return nullptr;
        }
        else
        {
            return root;
        }
    }
};