bumpsoo 2024-07-17 22:24:13 +09:00
parent 7f690e912e
commit ea63f9a6c6

View file

@ -0,0 +1,24 @@
# https://leetcode.com/problems/delete-nodes-and-return-forest
from typing import List, Optional
class TreeNode:
def __init__(self, val=0, left=None, right=None):
self.val = val
self.left = left
self.right = right
class Solution:
def delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
to_delete_set = set(to_delete)
ret = []
def dfs(node: Optional[TreeNode], is_root: bool) -> Optional[TreeNode]:
if not node:
return None
is_deleted = node.val in to_delete_set
if is_root and not is_deleted:
ret.append(node)
node.left = dfs(node.left, is_deleted)
node.right = dfs(node.right, is_deleted)
return None if is_deleted else node
dfs(root, True)
return ret