From dc9c3fd54d6787d8c9b11467baae4d3dfa2ee53e Mon Sep 17 00:00:00 2001
From: Dan Buch <dan@meatballhat.com>
Date: Sat, 21 Oct 2023 15:45:16 -0400
Subject: [PATCH] Binary tree right join goop

---
 leetcode/pyproject.toml |  1 -
 leetcode/stdlib.py      | 53 ++++++++++++++++++++++++++++++++++++++++-
 leetcode/stuff.py       | 44 ++++++++++++++++++++++++++++++++++
 leetcode/test_stuff.py  | 25 +++++++++++++++++++
 4 files changed, 121 insertions(+), 2 deletions(-)

diff --git a/leetcode/pyproject.toml b/leetcode/pyproject.toml
index 9de7a9e..de124fb 100644
--- a/leetcode/pyproject.toml
+++ b/leetcode/pyproject.toml
@@ -80,7 +80,6 @@ all = [
 [tool.black]
 target-version = ["py37"]
 line-length = 90
-skip-string-normalization = true
 
 [tool.ruff]
 target-version = "py37"
diff --git a/leetcode/stdlib.py b/leetcode/stdlib.py
index cdd4e9e..cafdb8d 100644
--- a/leetcode/stdlib.py
+++ b/leetcode/stdlib.py
@@ -4,6 +4,57 @@ import typing
 class ListNode:
     """ListNode is the leetcode "standard library" type used in linked lists"""
 
-    def __init__(self, val=0, next: typing.Optional["ListNode"] = None):  # no qa
+    def __init__(self, val: int = 0, next: typing.Optional["ListNode"] = None):  # no qa
         self.val = val
         self.next = next
+
+
+class Node:
+    """Node is the leetcode "standard library" type used in binary trees"""
+
+    def __init__(
+        self,
+        val: int = 0,
+        left: typing.Optional["Node"] = None,
+        right: typing.Optional["Node"] = None,
+        next: typing.Optional["Node"] = None,  # no qa
+    ):
+        self.val = val
+        self.left = left
+        self.right = right
+        self.next = next
+
+    # __repr__ was added by me
+    def __repr__(self) -> str:
+        filtered_parts = []
+
+        for key, value in [
+            ("val", self.val),
+            ("right", self.right),
+            ("left", self.left),
+            ("next", self.next),
+        ]:
+            if value is not None:
+                filtered_parts.append((key, value))
+
+        middle = ", ".join([f"{k}={v!r}" for k, v in filtered_parts])
+
+        return f"Node({middle})"
+
+    # __eq__ was added by me
+    def __eq__(self, other: "Node") -> bool:
+        return (
+            other is not None
+            and self.val == other.val
+            and self.left == other.left
+            and self.right == other.right
+            and self.next == other.next
+        )
+
+    # __list__ was added by me
+    def __list__(self) -> list[int | None]:
+        ret = [self.val]
+        ret += self.right.__list__() if self.right is not None else [None]
+        ret += self.left.__list__() if self.left is not None else [None]
+        ret += self.next.__list__() if self.next is not None else [None]
+        return ret
diff --git a/leetcode/stuff.py b/leetcode/stuff.py
index 29096a4..9786ca4 100644
--- a/leetcode/stuff.py
+++ b/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)
diff --git a/leetcode/test_stuff.py b/leetcode/test_stuff.py
index 941392b..ecd7458 100644
--- a/leetcode/test_stuff.py
+++ b/leetcode/test_stuff.py
@@ -152,3 +152,28 @@ def test_sort_linked_list(head: stdlib.ListNode | None, expected: stdlib.ListNod
     assert stuff.linked_list_to_list(
         stuff.sort_linked_list(head)
     ) == stuff.linked_list_to_list(expected)
+
+
+@pytest.mark.parametrize(
+    ("root", "expected"),
+    [
+        (
+            stdlib.Node(
+                1,
+                left=stdlib.Node(2, left=stdlib.Node(4), right=stdlib.Node(5)),
+                right=stdlib.Node(3, right=stdlib.Node(7)),
+            ),
+            [1, None, 2, 3, None, 4, 5, 7, None],
+        ),
+    ],
+)
+def test_connect_binary_tree_right(
+    root: stdlib.Node | None, expected: list[int | None] | None
+):
+    if expected is None:
+        assert root is None
+        return
+
+    connected, serialized = stuff.connect_binary_tree_right(root)
+    assert connected is not None
+    assert serialized == expected