Update test cases and sample code for 2026.05.27 homework

notes-and-examples/2026.05.27/homework.py

@@ -60,8 +60,12 @@ def rotate_left(node: Node) -> Node:
     new_right_child = node.right.left
 
     node.right = new_right_child
-    node.parent = new_parent
+    if new_right_child is not None:
+        new_right_child.parent = node
+
     new_parent.left = node
+    node.parent = new_parent
+    new_parent.parent = original_parent
 
     if original_parent is not None:
         if original_parent.left == node:
@@ -81,8 +85,12 @@ def rotate_right(node: Node) -> Node:
     new_left_child = node.left.right
 
     node.left = new_left_child
-    node.parent = new_parent
+    if new_left_child is not None:
+        new_left_child.parent = node
+
     new_parent.right = node
+    node.parent = new_parent
+    new_parent.parent = original_parent
 
     if original_parent is not None:
         if original_parent.left == node:
@@ -115,48 +123,86 @@ def trees_equal(a: Node | None, b: Node | None) -> bool:
 
 
 if __name__ == "__main__":
+    # note that matching is done by value for simplicity here
 
     def test_rotate_right_example():
         # The worksheet example: right rotate on 10.
         root = Node(
             10,
-            left=Node(5, left=Node(3, left=Node(2)), right=Node(7)),
-            right=Node(15),
+            left=Node(
+                5,
+                left=Node(3, left=Node(2, parent=Node(3)), parent=Node(5)),
+                right=Node(7, parent=Node(5)),
+                parent=Node(10),
+            ),
+            right=Node(15, parent=Node(10)),
         )
         expected = Node(
             5,
-            left=Node(3, left=Node(2)),
-            right=Node(10, left=Node(7), right=Node(15)),
+            left=Node(3, left=Node(2, parent=Node(3)), parent=Node(5)),
+            right=Node(
+                10,
+                left=Node(7, parent=Node(10)),
+                right=Node(15, parent=Node(10)),
+                parent=Node(5),
+            ),
         )
         return rotate_right(root), expected
 
+
     def test_rotate_left_example():
         # The worksheet example: left rotate on 16.
-        root = Node(16, right=Node(33, left=Node(22), right=Node(55)))
-        expected = Node(33, left=Node(16, right=Node(22)), right=Node(55))
+        root = Node(
+            16,
+            right=Node(
+                33,
+                left=Node(22, parent=Node(33)),
+                right=Node(55, parent=Node(33)),
+                parent=Node(16),
+            ),
+        )
+        expected = Node(
+            33,
+            left=Node(16, right=Node(22, parent=Node(16)), parent=Node(33)),
+            right=Node(55, parent=Node(33)),
+        )
         return rotate_left(root), expected
 
+
     def test_rotate_right_two_nodes():
         # Smallest case: a root with only a left child.
-        return rotate_right(Node(2, left=Node(1))), Node(1, right=Node(2))
+        return rotate_right(Node(2, left=Node(1, parent=Node(2)))), Node(1, right=Node(2, parent=Node(1)))
+
 
     def test_rotate_left_two_nodes():
         # Smallest case: a root with only a right child.
-        return rotate_left(Node(1, right=Node(2))), Node(2, left=Node(1))
+        return rotate_left(Node(1, right=Node(2, parent=Node(1)))), Node(2, left=Node(1, parent=Node(2)))
+
 
     def test_round_trip():
         # A left rotate undone by a right rotate restores the original tree.
-        original = Node(1, left=Node(0), right=Node(3, left=Node(2), right=Node(4)))
-        expected = Node(1, left=Node(0), right=Node(3, left=Node(2), right=Node(4)))
+        original = Node(1, left=Node(0, parent=Node(1)),
+                        right=Node(3, left=Node(2, parent=Node(3)), right=Node(4, parent=Node(3)), parent=Node(1)))
+        expected = Node(1, left=Node(0, parent=Node(1)),
+                        right=Node(3, left=Node(2, parent=Node(3)), right=Node(4, parent=Node(3)), parent=Node(1)))
         return rotate_right(rotate_left(original)), expected
 
+
     def test_rotate_with_parent():
         # run it on 33 instead of the root
         root = Node(16, right=Node(33, left=Node(22), right=Node(55)))
+        root.right.parent = root
+        root.right.left.parent = root.right
+        root.right.right.parent = root.right
         original = root.right
-        expected = Node(55, left=Node(33, left=Node(22)))
+        expected = Node(
+            55,
+            left=Node(33, left=Node(22, parent=Node(33)), parent=Node(55)),
+            parent=Node(16),
+        )
         return rotate_left(original), expected  # pyright: ignore[reportArgumentType]
 
+
     tests = [
         ("rotate_right on worksheet example (10)", test_rotate_right_example),
         ("rotate_left on worksheet example (16)", test_rotate_left_example),