Leetcode # 1110. Delete Nodes And Return Forest

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

Solution

class Solution:
  def delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
    to_delete = set(to_delete)
    # Depth-First Search
    stack = [(None, root, -1)]
    roots = {root}
    while stack:
      pre, cur, side = stack.pop()
      # if this node should be deleted
      if cur.val in to_delete:
        # delete the link between this and parent
        if pre:
          if side == 0:
            pre.left = None
          else:
            pre.right = None

        # delete this from roots if it's root
        if cur in roots:
          roots.remove(cur)
        # and add its children into roots
        if cur.left:
          roots.add(cur.left)
        if cur.right:
          roots.add(cur.right)

      # push into stack
      if cur.left:
        stack.append((cur, cur.left, 0))
      if cur.right:
        stack.append((cur, cur.right, 1))
    
    return roots

Last Updated on 2023/08/16 by A1go