Binary tree right join goop

leetcode/stuff.py

@@ -1,3 +1,4 @@
+import copy
 import typing
 
 import stdlib
@@ -199,3 +200,46 @@ def sort_linked_list(head: stdlib.ListNode | None) -> stdlib.ListNode | None:
         cur.next = None
 
     return ret
+
+
+def connect_binary_tree_right(
+    root: stdlib.Node | None,
+) -> tuple[stdlib.Node | None, list[int | None]]:
+    if root is None:
+        return None, []
+
+    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]):
+        for i in range(len(nodes)):
+            serialized.append(nodes[i].val)
+
+            if len(nodes) < i:
+                nodes[i].next = nodes[i + 1]
+
+        serialized.append(None)
+
+    return root, serialized
+
+
+def binary_tree_by_level(root: stdlib.Node) -> dict[int, list[stdlib.Node]]:
+    combined: dict[int, list[stdlib.Node]] = {}
+
+    for path in collect_binary_tree_levels(0, root):
+        level, node = path
+        combined.setdefault(level, [])
+        combined[level].insert(0, node)
+
+    return combined
+
+
+def collect_binary_tree_levels(
+    level: int, node: stdlib.Node | None
+) -> typing.Iterator[tuple[int, stdlib.Node]]:
+    if node is None:
+        return
+
+    yield (level, node)
+    yield from collect_binary_tree_levels(level + 1, node.right)
+    yield from collect_binary_tree_levels(level + 1, node.left)