diff --git a/src/binary_tree/dfs/path_sum_iii/__init__.py b/src/binary_tree/dfs/path_sum_iii/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/src/binary_tree/dfs/path_sum_iii/solution.py b/src/binary_tree/dfs/path_sum_iii/solution.py new file mode 100644 index 0000000..98a299b --- /dev/null +++ b/src/binary_tree/dfs/path_sum_iii/solution.py @@ -0,0 +1,25 @@ +from typing import Optional +from collections import defaultdict +from structures import TreeNode + + +class Solution: + def dfs(self, root: Optional[TreeNode], current_sum: int, target_sum: int): + if root is None: + return 0 + + current_sum += root.val + count: int = self.cache[current_sum - target_sum] + self.cache[current_sum] += 1 + + count += self.dfs(root.left, current_sum, target_sum) + count += self.dfs(root.right, current_sum, target_sum) + + self.cache[current_sum] -= 1 + + return count + + def pathSum(self, root: Optional[TreeNode], targetSum: int) -> int: + self.cache = defaultdict(int) + self.cache[0] = 1 + return self.dfs(root, 0, targetSum) diff --git a/structures/tree.py b/structures/tree.py index 6a976a2..6db4d3a 100644 --- a/structures/tree.py +++ b/structures/tree.py @@ -1,5 +1,10 @@ class TreeNode: - def __init__(self, val: int = 0, left: int | None = None, right: int | None = None): + def __init__( + self, + val: int = 0, + left: "TreeNode | None" = None, + right: "TreeNode | None" = None, + ): self.val = val self.left = left self.right = right diff --git a/tests/test_path_sum_iii.py b/tests/test_path_sum_iii.py new file mode 100644 index 0000000..2c03b9e --- /dev/null +++ b/tests/test_path_sum_iii.py @@ -0,0 +1,16 @@ +import pytest +from src.binary_tree.dfs.path_sum_iii.solution import Solution +from test_utils import TreeBuilder + + +@pytest.mark.parametrize( + "root, target_sum, expected", + [ + ([10, 5, -3, 3, 2, None, 11, 3, -2, None, 1], 8, 3), + ([5, 4, 8, 11, None, 13, 4, 7, 2, None, None, 5, 1], 22, 3), + ], +) +def test_leaf_similar(root: list, target_sum: int, expected: int): + root = TreeBuilder.from_list(root) + solution = Solution() + assert solution.pathSum(root, target_sum) == expected