230 lines
6.6 KiB
Python
230 lines
6.6 KiB
Python
# Implement left rotates and right rotates for trees.
|
|
#
|
|
# Example right rotate on 10:
|
|
#
|
|
# 10
|
|
# / \
|
|
# 5 15
|
|
# / \
|
|
# 3 7
|
|
# /
|
|
# 2
|
|
#
|
|
# Turns into:
|
|
#
|
|
# 5
|
|
# / \
|
|
# 3 10
|
|
# / / \
|
|
# 2 7 15
|
|
#
|
|
# Notice how the right child of 10's left child (7) becomes the left child of 10.
|
|
#
|
|
# Left rotates are the reverse. See this example on 16:
|
|
# 16
|
|
# \
|
|
# 33
|
|
# / \
|
|
# 22 55
|
|
#
|
|
# Turns into:
|
|
#
|
|
# 33
|
|
# / \
|
|
# 16 55
|
|
# \
|
|
# 22
|
|
#
|
|
# where the left child of 16's right child (22) turns into 16's right child.
|
|
#
|
|
# Copy this entire file, implement rotate_left and rotate_right,
|
|
# and run the tests using `python3 <name-of-file>.py`.
|
|
|
|
|
|
class Node:
|
|
def __init__(self, value, left=None, right=None, parent=None):
|
|
self.value = value
|
|
self.left: Node | None = left
|
|
self.right: Node | None = right
|
|
|
|
# required to re-set the child later
|
|
self.parent: Node | None = parent
|
|
|
|
|
|
def rotate_left(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(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
|
|
|
|
|
|
# --- Tests ---------------------------------------------------------------
|
|
# Both rotate functions are expected to return the new root of the rotated
|
|
# subtree (the parent would then point its child link at that new root).
|
|
|
|
|
|
def to_tuple(node: Node | None):
|
|
"""Turn a tree into nested (value, left, right) tuples for easy comparison."""
|
|
if node is None:
|
|
return None
|
|
return (
|
|
node.value,
|
|
to_tuple(node.left),
|
|
to_tuple(node.right),
|
|
node.parent.value if node.parent is not None else None,
|
|
)
|
|
|
|
|
|
def trees_equal(a: Node | None, b: Node | None) -> bool:
|
|
return to_tuple(a) == to_tuple(b)
|
|
|
|
|
|
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, 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, 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, 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, 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, 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, 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, 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),
|
|
("rotate_right on two nodes", test_rotate_right_two_nodes),
|
|
("rotate_left on two nodes", test_rotate_left_two_nodes),
|
|
("rotate_left then rotate_right restores the tree", test_round_trip),
|
|
("rotate_left works when there is a parent", test_rotate_with_parent),
|
|
]
|
|
|
|
passed = 0
|
|
for name, test in tests:
|
|
try:
|
|
result, expected = test()
|
|
if trees_equal(result, expected):
|
|
print(f"PASS: {name}")
|
|
passed += 1
|
|
else:
|
|
print(f"FAIL: {name}")
|
|
print(f" expected {to_tuple(expected)}")
|
|
print(f" got {to_tuple(result)}")
|
|
except Exception as e: # noqa: BLE001 - surface any bug as a failed test
|
|
print(f"ERROR: {name}: {type(e).__name__}: {e}")
|
|
|
|
print(f"\n{passed}/{len(tests)} tests passed")
|