Skip to content

1110. Delete Nodes And Return Forest 👍

  • Time: $O(n)$
  • Space: $O(\texttt{h} + |\texttt{toDelete}|)$
 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
 public:
  vector<TreeNode*> delNodes(TreeNode* root, vector<int>& to_delete) {
    vector<TreeNode*> ans;
    dfs(root, {begin(to_delete), end(to_delete)}, true, ans);
    return ans;
  }

 private:
  TreeNode* dfs(TreeNode*& root, const unordered_set<int>&& toDeleteSet,
                bool isRoot, vector<TreeNode*>& ans) {
    if (root == nullptr)
      return nullptr;

    const bool deleted = toDeleteSet.count(root->val);
    if (isRoot && !deleted)
      ans.push_back(root);

    // If root is deleted, both children have the possibility to be a new root
    root->left = dfs(root->left, move(toDeleteSet), deleted, ans);
    root->right = dfs(root->right, move(toDeleteSet), deleted, ans);
    return deleted ? nullptr : root;
  }
};

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
class Solution {
  public List<TreeNode> delNodes(TreeNode root, int[] to_delete) {
    List<TreeNode> ans = new ArrayList<>();
    Set<Integer> toDeleteSet = Arrays.stream(to_delete).boxed().collect(Collectors.toSet());

    dfs(root, toDeleteSet, true, ans);
    return ans;
  }

  private TreeNode dfs(TreeNode root, Set<Integer> toDeleteSet, boolean isRoot,
                       List<TreeNode> ans) {
    if (root == null)
      return null;

    final boolean deleted = toDeleteSet.contains(root.val);
    if (isRoot && !deleted)
      ans.add(root);

    // If root is deleted, both children have the possibility to be a new root
    root.left = dfs(root.left, toDeleteSet, deleted, ans);
    root.right = dfs(root.right, toDeleteSet, deleted, ans);
    return deleted ? null : root;
  }
}

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
class Solution:
  def delNodes(self, root: TreeNode, to_delete: List[int]) -> List[TreeNode]:
    ans = []
    toDeleteSet = set(to_delete)

    def dfs(root: TreeNode, isRoot: bool) -> TreeNode:
      if not root:
        return None

      deleted = root.val in toDeleteSet
      if isRoot and not deleted:
        ans.append(root)

      # If root is deleted, both children have the possibility to be a new root
      root.left = dfs(root.left, deleted)
      root.right = dfs(root.right, deleted)
      return None if deleted else root

    dfs(root, True)
    return ans