Sum binary tree path ints
This commit is contained in:
parent
e92d117a41
commit
cee338520e
@ -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
|
||||
|
@ -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
|
||||
|
Loading…
Reference in New Issue
Block a user