diff --git a/leetcode/stuff.py b/leetcode/stuff.py index a8a776f..a6b5cd7 100644 --- a/leetcode/stuff.py +++ b/leetcode/stuff.py @@ -251,3 +251,31 @@ def collect_binary_tree_levels( yield (level, node) yield from collect_binary_tree_levels(level + 1, node.right) yield from collect_binary_tree_levels(level + 1, node.left) + + +def sum_binary_tree_path_ints(root: stdlib.Node | None) -> int: + path_ints: list[int] = [] + + for path in collect_binary_tree_paths(root): + path_ints.append(int("".join([str(node.val) for node in path]))) + + return sum(path_ints) + + +def collect_binary_tree_paths( + node: stdlib.Node | None, +) -> typing.Iterator[list[stdlib.Node]]: + if node is None: + return + + if node.right is None and node.left is None: + yield [node] + return + + if node.right is not None: + for path in collect_binary_tree_paths(node.right): + yield [node] + path + + if node.left is not None: + for path in collect_binary_tree_paths(node.left): + yield [node] + path diff --git a/leetcode/test_stuff.py b/leetcode/test_stuff.py index ecd7458..6954a65 100644 --- a/leetcode/test_stuff.py +++ b/leetcode/test_stuff.py @@ -177,3 +177,20 @@ def test_connect_binary_tree_right( connected, serialized = stuff.connect_binary_tree_right(root) assert connected is not None assert serialized == expected + + +@pytest.mark.parametrize( + ("root", "expected"), + [ + ( + stdlib.Node( + 4, + right=stdlib.Node(0), + left=stdlib.Node(9, right=stdlib.Node(1), left=stdlib.Node(5)), + ), + 1026, + ), + ], +) +def test_connect_binary_tree_sum_numbers(root: stdlib.Node | None, expected: int): + assert stuff.sum_binary_tree_path_ints(root) == expected