From e92d117a41cd10c7f7c4d030d36be430b518230b Mon Sep 17 00:00:00 2001 From: Dan Buch Date: Sat, 21 Oct 2023 16:19:13 -0400 Subject: [PATCH] Return the correct right-connected binary tree root --- leetcode/stuff.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/leetcode/stuff.py b/leetcode/stuff.py index 9786ca4..a8a776f 100644 --- a/leetcode/stuff.py +++ b/leetcode/stuff.py @@ -211,16 +211,24 @@ def connect_binary_tree_right( by_level = binary_tree_by_level(copy.deepcopy(root)) serialized: list[int | None] = [] - for _, nodes in sorted(by_level.items(), key=lambda p: p[0]): + print("") + + if 0 not in by_level or len(by_level[0]) == 0: + return None, [] + + connected_root = by_level[0][0] + + for level, nodes in sorted(by_level.items(), key=lambda p: p[0]): for i in range(len(nodes)): serialized.append(nodes[i].val) - if len(nodes) < i: + if len(nodes) > i + 1: + print(f"{'-' * level}> connecting {nodes[i].val} -> {nodes[i + 1].val}") nodes[i].next = nodes[i + 1] serialized.append(None) - return root, serialized + return connected_root, serialized def binary_tree_by_level(root: stdlib.Node) -> dict[int, list[stdlib.Node]]: