Delete nodes and return forest

Author: volodimyr

1110.delete_nodes_and_return_forest/delete.py

@@ -0,0 +1,69 @@
+# 1110. Delete nodes and return forest
+# Topics: 'Array', 'Hash Table', 'Tree', 'Depth-First Search', 'Binary Tree'
+# Level: 'Medium'
+
+# Given the root of a binary tree, each node in the tree has a distinct value.
+
+# After deleting all nodes with a value in to_delete, we are left with a forest (a disjoint union of trees).
+
+# Return the roots of the trees in the remaining forest. You may return the result in any order.
+
+ 
+
+# Example 1:
+
+# Input: root = [1,2,3,4,5,6,7], to_delete = [3,5]
+# Output: [[1,2,null,4],[6],[7]]
+
+# Example 2:
+
+# Input: root = [1,2,4,null,3], to_delete = [3]
+# Output: [[1,2,4]]
+
+ 
+
+# Constraints:
+
+#     The number of nodes in the given tree is at most 1000.
+#     Each node has a distinct value between 1 and 1000.
+#     to_delete.length <= 1000
+#     to_delete contains distinct values between 1 and 1000.
+
+# Definition for a binary tree node.
+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
+    
+    def postorder(self, root: Optional[TreeNode]):
+        if not root:
+            return
+        self.postorder(root.left)
+        self.postorder(root.right)
+        print(root.val, end=" ")
+    
+class Solution:
+    def delNodes(self, root: Optional[TreeNode], to_delete: List[int]) -> List[TreeNode]:
+        dl = set(to_delete)
+        res = []
+        def dfs(root: Optional[TreeNode]) -> Optional[TreeNode]:
+            if not root:
+                return root
+            root.left = dfs(root.left)
+            root.right = dfs(root.right)
+            if root.val in dl:
+                if root.left:
+                    res.append(root.left)
+                if root.right:
+                    res.append(root.right)
+                return None
+            return root
+            
+        dfs(root)
+        if root.val not in dl:
+            res.append(root)
+        return res

1110.delete_nodes_and_return_forest/test_delete.py

@@ -0,0 +1,163 @@
+import unittest
+from typing import Optional, List
+from delete import TreeNode, Solution
+
+class TestDeleteNodes(unittest.TestCase):
+    
+    def tree_to_list(self, root: Optional[TreeNode]) -> List:
+        """Helper: Convert tree to level-order list for comparison"""
+        if not root:
+            return []
+        result = []
+        queue = [root]
+        while queue:
+            node = queue.pop(0)
+            if node:
+                result.append(node.val)
+                queue.append(node.left)
+                queue.append(node.right)
+            else:
+                result.append(None)
+        # Remove trailing Nones
+        while result and result[-1] is None:
+            result.pop()
+        return result
+    
+    def test_example_1(self):
+        """Test case from example 1: delete nodes 3 and 5"""
+        # Tree: [1,2,3,4,5,6,7]
+        root = TreeNode(1,
+                       TreeNode(2, TreeNode(4), TreeNode(5)),
+                       TreeNode(3, TreeNode(6), TreeNode(7)))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [3, 5])
+        
+        # Convert results to lists for easier comparison
+        result_lists = sorted([self.tree_to_list(tree) for tree in result])
+        
+        # Expected: [[1,2,null,4], [6], [7]]
+        expected = sorted([[1, 2, None, 4], [6], [7]])
+        
+        self.assertEqual(result_lists, expected)
+    
+    def test_example_2(self):
+        """Test case from example 2: delete node 3"""
+        # Tree: [1,2,4,null,3]
+        root = TreeNode(1,
+                       TreeNode(2, None, TreeNode(3)),
+                       TreeNode(4))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [3])
+        
+        result_lists = [self.tree_to_list(tree) for tree in result]
+        expected = [[1, 2, 4]]
+        
+        self.assertEqual(result_lists, expected)
+    
+    def test_delete_root_only(self):
+        """Delete only the root node"""
+        root = TreeNode(1, TreeNode(2), TreeNode(3))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [1])
+        
+        result_lists = sorted([self.tree_to_list(tree) for tree in result])
+        expected = sorted([[2], [3]])
+        
+        self.assertEqual(result_lists, expected)
+    
+    def test_delete_nothing(self):
+        """Delete no nodes - return original tree"""
+        root = TreeNode(1, TreeNode(2), TreeNode(3))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [])
+        
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.tree_to_list(result[0]), [1, 2, 3])
+    
+    def test_delete_all_nodes(self):
+        """Delete all nodes in tree"""
+        root = TreeNode(1, TreeNode(2), TreeNode(3))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [1, 2, 3])
+        
+        self.assertEqual(result, [])
+    
+    def test_single_node_delete(self):
+        """Single node tree, delete it"""
+        root = TreeNode(1)
+        
+        solution = Solution()
+        result = solution.delNodes(root, [1])
+        
+        self.assertEqual(result, [])
+    
+    def test_single_node_keep(self):
+        """Single node tree, keep it"""
+        root = TreeNode(1)
+        
+        solution = Solution()
+        result = solution.delNodes(root, [])
+        
+        self.assertEqual(len(result), 1)
+        self.assertEqual(result[0].val, 1)
+    
+    def test_delete_leaf_nodes(self):
+        """Delete only leaf nodes"""
+        root = TreeNode(1,
+                       TreeNode(2, TreeNode(4), TreeNode(5)),
+                       TreeNode(3))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [4, 5, 3])
+        
+        self.assertEqual(len(result), 1)
+        self.assertEqual(self.tree_to_list(result[0]), [1, 2])
+    
+    def test_delete_internal_nodes(self):
+        """Delete internal nodes creating multiple forests"""
+        root = TreeNode(1,
+                       TreeNode(2, TreeNode(4), TreeNode(5)),
+                       TreeNode(3, TreeNode(6), TreeNode(7)))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [2, 3])
+        
+        result_lists = sorted([self.tree_to_list(tree) for tree in result])
+        expected = sorted([[1], [4], [5], [6], [7]])
+        
+        self.assertEqual(result_lists, expected)
+    
+    def test_linear_tree_delete_middle(self):
+        """Linear tree (like linked list), delete middle node"""
+        root = TreeNode(1, TreeNode(2, TreeNode(3, TreeNode(4))))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [2])
+        
+        result_lists = sorted([self.tree_to_list(tree) for tree in result])
+        expected = sorted([[1], [3, 4]])
+        
+        self.assertEqual(result_lists, expected)
+    
+    def test_delete_node_with_one_child(self):
+        """Delete node that has only one child"""
+        root = TreeNode(1,
+                       TreeNode(2, TreeNode(4)),
+                       TreeNode(3))
+        
+        solution = Solution()
+        result = solution.delNodes(root, [2])
+        
+        result_lists = sorted([self.tree_to_list(tree) for tree in result])
+        expected = sorted([[1, None, 3], [4]])
+        
+        self.assertEqual(result_lists, expected)
+
+
+if __name__ == '__main__':
+    unittest.main(verbosity=2)