1110

1110. Delete Nodes And Return Forest (Medium)

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 and 1000.
• to_delete.length <= 1000
• to_delete contains distinct values between 1 and 1000.

Related Topics:
Tree, Depth-first Search

Solution 1.

// 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;
    }
};