【Leetcode】1110. Delete Nodes And Return Forest

题目地址：

https://leetcode.com/problems/delete-nodes-and-return-forest/

给定一棵二叉树，再给定一个数组$A$，要求在二叉树中将$A$中存在的节点删掉，并返回得到的森林的每棵树的树根。

这个问题需要同时做两件事，一个是要把含$A$中数字的节点删掉，另一个是要将删完了之后产生的森林的树根加入列表中。删除节点的方法就是，在DFS的时候返回当前节点的子树做删除操作完成之后的树根，所以也就是说，如果当前节点就是要删的，则返回null，否则返回当前节点自己。而收集森林树根的过程必须先于返回节点执行，也就是这里应该是要后序遍历。但问题是，如何判断当前节点是否应该加入答案列表中呢。我们需要在DFS的时候再传一个参数，即当前节点是否是森林的树根。如果是树根，就加入。如果当前节点要删，显然就不能成为树根；否则的话，要么其没有父亲，要么其父亲会被删，这两种情况下当前节点会成为森林的某个树根，从而要加入答案。代码如下：

/**
 * 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;
    unordered_set<int> st(to_delete.begin(), to_delete.end());
    dfs(root, nullptr, st, res);
    return res;
  }

  TreeNode* dfs(TreeNode* p, TreeNode* pa, auto& st, auto& res) {
    if (!p) return nullptr;
    p->left = dfs(p->left, p, st, res);
    p->right = dfs(p->right, p, st, res);
    if (st.count(p->val)) return nullptr;
    if (!pa || st.count(pa->val)) res.push_back(p);
    return p;
  }
};

时间复杂度$O(n)$，空间$O(h)$。