From ea63f9a6c604c9e31fbbbb634205884e2de47437 Mon Sep 17 00:00:00 2001 From: bumpsoo Date: Wed, 17 Jul 2024 22:24:13 +0900 Subject: [PATCH] https://leetcode.com/problems/delete-nodes-and-return-forest --- medium/delete_nodes_and_return_forest.py | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) create mode 100644 medium/delete_nodes_and_return_forest.py 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