Implement a sample solution which passes all test cases

notes-and-examples/2026.07.01/red_black_tree.py

@@ -84,17 +84,95 @@ class RedBlackTree:
         elif root.value < value:
             self.delete(value, root.right, root)
 
-    def rebalance_from_just_inserted(self, node: Node):
-        # This is what you'll be implementing.
-        #
-        # By the time this is called, `node` has already been inserted like a
-        # normal BST node, colored red, and had its `parent` pointer set (see
-        # insert()). Your job is to restore the red-black properties by
-        # rebalancing the subtree around `node`.
-        #
-        # Rebalancing is only needed when node.parent is red. Remember to keep
-        # the root black at the end.
-        pass
+    def rebalance_from_just_inserted(self, node: Node | None):
+        if node is None:
+            return
+
+        parent = node.parent
+        if parent is not None and not parent.is_red:
+            return
+        elif parent is None:
+            node.is_red = False
+            return
+
+        grandparent = parent.parent
+        if grandparent is None:
+            return
+
+        uncle = grandparent.right if grandparent.left == parent else grandparent.left
+
+        if uncle is not None and uncle.is_red:
+            uncle.is_red = False
+            parent.is_red = False
+            grandparent.is_red = True
+
+            self.rebalance_from_just_inserted(grandparent)
+            return
+
+        if grandparent.left == parent and parent.right == node:
+            node, parent = parent, self.rotate_left(parent)
+        elif grandparent.right == parent and parent.left == node:
+            node, parent = parent, self.rotate_right(parent)
+
+        new_grandparent: Node | None = None
+        if grandparent.left == parent and parent.left == node:
+            new_grandparent = self.rotate_right(grandparent)
+        elif grandparent.right == parent and parent.right == node:
+            new_grandparent = self.rotate_left(grandparent)
+
+        if grandparent == self.root and new_grandparent is not None:
+            self.root = new_grandparent
+            self.root.is_red = False
+            grandparent.is_red = True
+
+
+    def rotate_left(self, node: Node) -> Node:
+        if node.right is None:
+            return node
+
+        original_parent = node.parent
+        new_parent = node.right
+        new_right_child = node.right.left
+
+        node.right = new_right_child
+        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:
+                original_parent.left = new_parent
+            elif original_parent.right == node:
+                original_parent.right = new_parent
+
+        return new_parent
+
+    def rotate_right(self, node: Node) -> Node:
+        if node.left is None:
+            return node
+
+        original_parent = node.parent
+        new_parent = node.left
+        new_left_child = node.left.right
+
+        node.left = new_left_child
+        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:
+                original_parent.left = new_parent
+            elif original_parent.right == node:
+                original_parent.right = new_parent
+
+        return new_parent
 
     def visualize(self) -> str:
         # renders the tree top-down with the root on top and branches