feature: path sum ii dfs

solairerove · Nov 16, 2023 · beb7af0 · beb7af0
1 parent 9d7ab29
commit beb7af0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 4 deletions.
diff --git a/trees/PathSumII.py b/trees/PathSumII.py
@@ -4,7 +4,10 @@
 
 
 # O(n) time || O(max(n, log(n)) space
-def path_sum(self, root: Optional[TreeNode], target_sum: int) -> List[List[int]]:
+def path_sum_rec(self, root: Optional[TreeNode], target_sum: int) -> List[List[int]]:
+    if not root:
+        return []
+
     res = []
 
     def dfs(node, curr_sum, curr_path):
@@ -25,3 +28,25 @@ def dfs(node, curr_sum, curr_path):
     dfs(root, 0, [])
 
     return res
+
+
+# O(n) time || O(max(n, log(n)) space
+def path_sum_dfs(self, root: Optional[TreeNode], target_sum: int) -> List[List[int]]:
+    if not root:
+        return []
+
+    res = []
+    stack = [(root, root.val, [root.val])]
+    while stack:
+        node, curr_sum, curr_path = stack.pop()
+
+        if not node.left and not node.right and curr_sum == target_sum:
+            res.append(curr_path)
+
+        if node.left:
+            stack.append((node.left, curr_sum + node.left.val, curr_path + [node.left.val]))
+
+        if node.right:
+            stack.append((node.right, curr_sum + node.right.val, curr_path + [node.right.val]))
+
+    return res
diff --git a/trees/test/test_path_sum.py b/trees/test/test_path_sum.py
@@ -1,6 +1,6 @@
 import unittest
 
-from trees.PathSumII import path_sum
+from trees.PathSumII import path_sum_rec, path_sum_dfs
 from trees.TreeNode import TreeNode
 
 
@@ -26,7 +26,8 @@ def test_path_sum(self):
                 )
             )
         )
-        self.assertEqual([[5, 4, 11, 2], [5, 8, 4, 5]], path_sum(self, root, 22))
+        self.assertEqual([[5, 4, 11, 2], [5, 8, 4, 5]], path_sum_rec(self, root, 22))
+        self.assertEqual([[5, 8, 4, 5], [5, 4, 11, 2]], path_sum_dfs(self, root, 22))
 
     def test_path_sum_1(self):
         root = TreeNode(
@@ -35,7 +36,8 @@ def test_path_sum_1(self):
             right=TreeNode(3)
         )
 
-        self.assertEqual([], path_sum(self, root, 0))
+        self.assertEqual([], path_sum_rec(self, root, 0))
+        self.assertEqual([], path_sum_dfs(self, root, 0))
 
 
 if __name__ == '__main__':