Formatted question description: https://leetcode.ca/all/1110.html

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/

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

Java

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
        if (root == null)
            return new ArrayList<TreeNode>();
        Set<Integer> deleteSet = new HashSet<Integer>();
        for (int value : to_delete)
            deleteSet.add(value);
        List<TreeNode> nodesList = new ArrayList<TreeNode>();
        Map<TreeNode, TreeNode> map = new HashMap<TreeNode, TreeNode>();
        Queue<TreeNode> queue = new LinkedList<TreeNode>();
        queue.offer(root);
        while (!queue.isEmpty()) {
            TreeNode node = queue.poll();
            nodesList.add(node);
            TreeNode left = node.left, right = node.right;
            if (left != null) {
                map.put(left, node);
                queue.offer(left);
            }
            if (right != null) {
                map.put(right, node);
                queue.offer(right);
            }
        }
        List<TreeNode> forest = new ArrayList<TreeNode>();
        if (!deleteSet.contains(root.val))
            forest.add(root);
        for (int i = nodesList.size() - 1; i >= 0 && !deleteSet.isEmpty(); i--) {
            TreeNode node = nodesList.get(i);
            if (deleteSet.contains(node.val)) {
                TreeNode parent = map.get(node);
                if (parent != null) {
                    if (node == parent.left)
                        parent.left = null;
                    else if (node == parent.right)
                        parent.right = null;
                }
                TreeNode left = node.left, right = node.right;
                if (left != null)
                    forest.add(left);
                if (right != null)
                    forest.add(right);
                deleteSet.remove(node.val);
            }
        }
        return forest;
    }
}

All Problems

All Solutions