Python/data_structures/binary_tree/splay_tree.py at 31ef93529f9c5612f1799f77c0875501166f774e · TheAlgorithms/Python · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
"""
Splay Tree implementation in Python.

A Splay Tree is a self-adjusting binary search tree where recently accessed
elements are moved closer to the root using rotations (splaying).
This improves access times for frequently used elements.

Reference: https://en.wikipedia.org/wiki/Splay_tree
Author: yeshuawm999
"""

from __future__ import annotations


class Node:
    """A node in the Splay Tree."""

    def __init__(
        self,
        key: int,
        parent: Node | None = None,
        left: Node | None = None,
        right: Node | None = None,
    ) -> None:
        self.key = key
        self.parent = parent
        self.left = left
        self.right = right


class SplayTree:
    """A self-adjusting Binary Search Tree (Splay Tree)."""

    def __init__(self) -> None:
        self.root: Node | None = None

    # --- Basic Rotation Operations ---

    def _rotate_left(self, x: Node) -> None:
        """Perform a left rotation around node x."""
        y = x.right
        if not y:
            return
        x.right = y.left
        if y.left:
            y.left.parent = x

        y.parent = x.parent
        if not x.parent:
            self.root = y
        elif x == x.parent.left:
            x.parent.left = y
        else:
            x.parent.right = y

        y.left = x
        x.parent = y

    def _rotate_right(self, x: Node) -> None:
        """Perform a right rotation around node x."""
        y = x.left
        if not y:
            return
        x.left = y.right
        if y.right:
            y.right.parent = x

        y.parent = x.parent
        if not x.parent:
            self.root = y
        elif x == x.parent.right:
            x.parent.right = y
        else:
            x.parent.left = y

        y.right = x
        x.parent = y

    # --- Core Splay Operation ---

    def _splay(self, x: Node) -> None:
        """Moves node x to the root of the tree using rotations."""
        while x.parent:
            parent = x.parent
            grandparent = parent.parent

            if not grandparent:
                # Zig Case (x is a child of the root)
                # One single rotation:
                # Right if x is a left child, Left if x is a right child
                if x == parent.left:
                    self._rotate_right(parent)
                else:
                    self._rotate_left(parent)
            elif x == parent.left and parent == grandparent.left:
                # Zig-Zig (both left)
                self._rotate_right(grandparent)
                self._rotate_right(parent)
            elif x == parent.right and parent == grandparent.right:
                # Zig-Zig (both right)
                self._rotate_left(grandparent)
                self._rotate_left(parent)
            elif x == parent.left and parent == grandparent.right:
                # Zig-Zag (left-right)
                self._rotate_right(parent)
                self._rotate_left(grandparent)
            else:
                # Zig-Zag (right-left)
                self._rotate_left(parent)
                self._rotate_right(grandparent)

    # --- Search Method ---

    def search(self, key: int) -> bool:
        """Search for a key. If found, splay it to the root."""
        current = self.root
        found_node = None
        while current:
            if key == current.key:
                found_node = current
                break
            current = current.left if key < current.key else current.right

        if found_node:
            self._splay(found_node)
            return True
        return False


if __name__ == "__main__":
    """
    Example run:
    >>> tree = SplayTree()
    >>> tree.root = Node(10)
    >>> tree.root.left = Node(5, parent=tree.root)
    >>> tree.root.right = Node(15, parent=tree.root)
    >>> print("Before search:", tree.root.key)
    >>> found = tree.search(5)
    >>> print("Found:", found)
    >>> print("After splay, new root:", tree.root.key)
    """

    tree = SplayTree()
    tree.root = Node(10)
    tree.root.left = Node(5, parent=tree.root)
    tree.root.right = Node(15, parent=tree.root)

    print("Before search:", tree.root.key)
    found = tree.search(5)
    print("Found:", found)
    print("After splay, new root:", tree.root.key)