diff --git a/medium/delete_nodes_and_return_forest.py b/medium/delete_nodes_and_return_forest.py new file mode 100644 index 0000000..382dd72 --- /dev/null +++ b/medium/delete_nodes_and_return_forest.py @@ -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