Return the correct right-connected binary tree root

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]]: