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
and1000
. to_delete.length <= 1000
to_delete
contains distinct values between1
and1000
.
Related Topics:
Tree, Depth-first Search
// OJ: https://leetcode.com/problems/delete-nodes-and-return-forest/
// Author: github.com/lzl124631x
// Time: O(N)
// Space: O(D + H)
class Solution {
vector<TreeNode*> ans;
unordered_set<int> s;
void dfs(TreeNode *root, TreeNode *p = NULL) {
if (!root) return;
dfs(root->left, root);
dfs(root->right, root);
if (s.count(root->val)) {
if (p != NULL) {
if (p->left == root) p->left = NULL;
else p->right = NULL;
}
if (root->left) ans.push_back(root->left);
if (root->right) ans.push_back(root->right);
} else if (p == NULL) ans.push_back(root);
}
public:
vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
s = unordered_set<int>(begin(to_delete), end(to_delete));
dfs(root);
return ans;
}
};