diff --git a/data_structures/binary_tree/splay_tree.py b/data_structures/binary_tree/splay_tree.py new file mode 100644 index 000000000000..358706cc3a16 --- /dev/null +++ b/data_structures/binary_tree/splay_tree.py @@ -0,0 +1,334 @@ +""" +Splay Tree implementation - A self-adjusting binary search tree. + +A Splay tree is a self-adjusting binary search tree where recently accessed +elements are moved to the root through rotations. This provides amortized +O(log n) time complexity for search, insert, and delete operations. + +The splaying operation moves a node to the root by performing a series of +rotations, making frequently accessed elements faster to access in the future. + +For more information, see: https://en.wikipedia.org/wiki/Splay_tree + +Time Complexity (amortized): +- Search: O(log n) +- Insert: O(log n) +- Delete: O(log n) + +Space Complexity: O(n) + +Operations: +- Zig: Single rotation (when parent is root) +- Zig-zig: Double rotation in same direction +- Zig-zag: Double rotation in opposite directions + +Example: +>>> tree = SplayTree() +>>> _ = tree.insert(10) +>>> _ = tree.insert(20) +>>> _ = tree.insert(30) +>>> _ = tree.insert(5) +>>> _ = tree.insert(15) +>>> list(tree.inorder()) +[5, 10, 15, 20, 30] +>>> tree.search(15) +True +>>> tree.search(25) +False +>>> _ = tree.delete(20) +>>> list(tree.inorder()) +[5, 10, 15, 30] +""" + +from __future__ import annotations + +from collections.abc import Iterator +from dataclasses import dataclass +from typing import Any, Self + + +@dataclass +class SplayNode: + """A node in the Splay Tree.""" + + value: Any + left: SplayNode | None = None + right: SplayNode | None = None + parent: SplayNode | None = None + + def __iter__(self) -> Iterator[Any]: + """Inorder traversal iterator.""" + yield from self.left or [] + yield self.value + yield from self.right or [] + + def __repr__(self) -> str: + from pprint import pformat + + if self.left is None and self.right is None: + return str(self.value) + return pformat({f"{self.value}": (self.left, self.right)}, indent=1) + + @property + def is_right(self) -> bool: + """Check if this node is the right child of its parent.""" + return bool(self.parent and self is self.parent.right) + + +@dataclass +class SplayTree: + """ + Splay Tree implementation - A self-adjusting BST. + + This tree automatically moves recently accessed elements to the root + through rotations, providing amortized O(log n) performance for all operations. + """ + + root: SplayNode | None = None + + def __bool__(self) -> bool: + """Return True if the tree is not empty.""" + return self.root is not None + + def __iter__(self) -> Iterator[Any]: + """Iterate over the tree in inorder traversal.""" + yield from self.root or [] + + def __len__(self) -> int: + """Return the number of nodes in the tree.""" + return sum(1 for _ in self) + + def __str__(self) -> str: + """Return a string representation of the tree.""" + return str(self.root) + + def _rotate_right(self, node: SplayNode) -> None: + """Perform a right rotation on the given node.""" + if not node.left: + return + + left_child = node.left + node.left = left_child.right + + if left_child.right: + left_child.right.parent = node + + left_child.parent = node.parent + + if not node.parent: + self.root = left_child + elif node is node.parent.right: + node.parent.right = left_child + else: + node.parent.left = left_child + + left_child.right = node + node.parent = left_child + + def _rotate_left(self, node: SplayNode) -> None: + """Perform a left rotation on the given node.""" + if not node.right: + return + + right_child = node.right + node.right = right_child.left + + if right_child.left: + right_child.left.parent = node + + right_child.parent = node.parent + + if not node.parent: + self.root = right_child + elif node is node.parent.left: + node.parent.left = right_child + else: + node.parent.right = right_child + + right_child.left = node + node.parent = right_child + + def _splay(self, node: SplayNode) -> None: + """ + Splay the given node to the root through a series of rotations. + + The splaying operation uses three types of rotations: + - Zig: Single rotation when parent is root + - Zig-zig: Double rotation in same direction + - Zig-zag: Double rotation in opposite directions + """ + while node.parent: + parent = node.parent + grandparent = parent.parent + + if not grandparent: + # Zig case: parent is root + if node is parent.left: + self._rotate_right(parent) + else: + self._rotate_left(parent) + elif (node is parent.left and parent is grandparent.left) or ( + node is parent.right and parent is grandparent.right + ): + # Zig-zig case: same direction + if parent is grandparent.left: + self._rotate_right(grandparent) + self._rotate_right(parent) + else: + self._rotate_left(grandparent) + self._rotate_left(parent) + elif node is parent.left: + # Zig-zag case: opposite directions (left child) + self._rotate_right(parent) + self._rotate_left(grandparent) + else: + # Zig-zag case: opposite directions (right child) + self._rotate_left(parent) + self._rotate_right(grandparent) + + def _find_node(self, value: Any) -> SplayNode | None: + """Find a node with the given value, splaying it to root if found.""" + current = self.root + while current: + if value == current.value: + self._splay(current) + return current + elif value < current.value: + if not current.left: + break + current = current.left + else: + if not current.right: + break + current = current.right + + # If we found a node (even if not exact match), splay it + if current: + self._splay(current) + return None + + def search(self, value: Any) -> bool: + """Search for a value in the tree. Returns True if found.""" + return self._find_node(value) is not None + + def insert(self, value: Any) -> Self: + """Insert a value into the splay tree.""" + if not self.root: + self.root = SplayNode(value) + return self + + # Find the insertion point + current = self.root + while True: + if value < current.value: + if current.left: + current = current.left + else: + current.left = SplayNode(value, parent=current) + self._splay(current.left) + break + elif value > current.value: + if current.right: + current = current.right + else: + current.right = SplayNode(value, parent=current) + self._splay(current.right) + break + else: + # Value already exists, splay the existing node + self._splay(current) + break + + return self + + def delete(self, value: Any) -> Self: + """Delete a value from the splay tree.""" + node = self._find_node(value) + if not node: + return self + + # Node to delete is now at root + left_tree = node.left + right_tree = node.right + + # Remove the root + if left_tree: + left_tree.parent = None + if right_tree: + right_tree.parent = None + + # If no left subtree, right becomes root + if not left_tree: + self.root = right_tree + return self + + # Find the maximum in left subtree + max_left = left_tree + while max_left.right: + max_left = max_left.right + + # Splay the maximum to root of left subtree + self._splay(max_left) + + # Attach right subtree to the new root + max_left.right = right_tree + if right_tree: + right_tree.parent = max_left + + self.root = max_left + return self + + def inorder(self) -> Iterator[Any]: + """Return an inorder iterator.""" + yield from self + + def preorder(self) -> Iterator[Any]: + """Return a preorder iterator.""" + + def _preorder(node: SplayNode | None) -> Iterator[Any]: + if node: + yield node.value + yield from _preorder(node.left) + yield from _preorder(node.right) + + yield from _preorder(self.root) + + def postorder(self) -> Iterator[Any]: + """Return a postorder iterator.""" + + def _postorder(node: SplayNode | None) -> Iterator[Any]: + if node: + yield from _postorder(node.left) + yield from _postorder(node.right) + yield node.value + + yield from _postorder(self.root) + + def get_min(self) -> Any: + """Get the minimum value in the tree.""" + if not self.root: + raise ValueError("Tree is empty") + current = self.root + while current.left: + current = current.left + self._splay(current) + return current.value + + def get_max(self) -> Any: + """Get the maximum value in the tree.""" + if not self.root: + raise ValueError("Tree is empty") + current = self.root + while current.right: + current = current.right + self._splay(current) + return current.value + + def is_empty(self) -> bool: + """Check if the tree is empty.""" + return self.root is None + + def clear(self) -> Self: + """Clear the tree.""" + self.root = None + return self