Skip to content

Tree Data Structures

Hierarchical data structures optimized for search and retrieval.

Binary Search Tree

Bases: BinaryTree[T]

A strictly typed, memory-optimized Binary Search Tree.

Source code in pure_python_ds/trees/binary_search_tree.py
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
class BinarySearchTree(BinaryTree[T]):
    """A strictly typed, memory-optimized Binary Search Tree."""

    def insert(self, value: T) -> None:
        """Iterative O(log n) insertion to maintain strict order."""
        new_node = TreeNode(value)
        if not self.root:
            self.root = new_node
            return

        current = self.root
        while True:
            if value < current.value:
                if not current.left:
                    current.left = new_node
                    return
                current = current.left
            elif value > current.value:
                if not current.right:
                    current.right = new_node
                    return
                current = current.right
            else:
                # Value already exists; ignoring duplicates for strict BST rules
                return

    def find(self, value: T) -> bool:
        """Iterative O(log n) search returning True if value exists."""
        current = self.root
        while current:
            if value == current.value:
                return True
            elif value < current.value:
                current = current.left
            else:
                current = current.right
        return False

    def search(self, key: T) -> Optional[T]:
        """Searches for a key in the BST and returns the value if found."""
        return self._search_recursive(self.root, key)

    def _search_recursive(self, node, key):
        if node is None or node.value == key:
            return node

        if key < node.value:
            return self._search_recursive(node.left, key)
        return self._search_recursive(node.right, key)

    def delete(self, key: T) -> None:
        """Removes a key from the BST. Targets 100% coverage for deletion logic."""
        self.root = self._delete_recursive(self.root, key)

    def _delete_recursive(
        self, node: Optional[TreeNode[T]], key: T
    ) -> Optional[TreeNode[T]]:
        if not node:
            return None

        if key < node.value:
            node.left = self._delete_recursive(node.left, key)
        elif key > node.value:
            node.right = self._delete_recursive(node.right, key)
        else:
            # Case 1 & 2: Node with only one child or no child
            if not node.left:
                return node.right
            elif not node.right:
                return node.left

            # Case 3: Node with two children
            # Get the inorder successor (smallest in the right subtree)
            successor = self._min_value_node(node.right)
            node.value = successor.value
            # Delete the inorder successor
            node.right = self._delete_recursive(node.right, successor.value)

        return node

    def _min_value_node(self, node: TreeNode[T]) -> TreeNode[T]:
        current = node
        while current.left:
            current = current.left
        return current

delete(key)

Removes a key from the BST. Targets 100% coverage for deletion logic.

Source code in pure_python_ds/trees/binary_search_tree.py
63
64
65
def delete(self, key: T) -> None:
    """Removes a key from the BST. Targets 100% coverage for deletion logic."""
    self.root = self._delete_recursive(self.root, key)

find(value)

Iterative O(log n) search returning True if value exists.

Source code in pure_python_ds/trees/binary_search_tree.py
39
40
41
42
43
44
45
46
47
48
49
def find(self, value: T) -> bool:
    """Iterative O(log n) search returning True if value exists."""
    current = self.root
    while current:
        if value == current.value:
            return True
        elif value < current.value:
            current = current.left
        else:
            current = current.right
    return False

insert(value)

Iterative O(log n) insertion to maintain strict order.

Source code in pure_python_ds/trees/binary_search_tree.py
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
def insert(self, value: T) -> None:
    """Iterative O(log n) insertion to maintain strict order."""
    new_node = TreeNode(value)
    if not self.root:
        self.root = new_node
        return

    current = self.root
    while True:
        if value < current.value:
            if not current.left:
                current.left = new_node
                return
            current = current.left
        elif value > current.value:
            if not current.right:
                current.right = new_node
                return
            current = current.right
        else:
            # Value already exists; ignoring duplicates for strict BST rules
            return

search(key)

Searches for a key in the BST and returns the value if found.

Source code in pure_python_ds/trees/binary_search_tree.py
51
52
53
def search(self, key: T) -> Optional[T]:
    """Searches for a key in the BST and returns the value if found."""
    return self._search_recursive(self.root, key)

AVL Tree

Bases: BinarySearchTree[T]

A strictly typed, self-balancing AVL Tree guaranteeing O(log n) operations.

Source code in pure_python_ds/trees/avl_tree.py
 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
152
153
154
155
156
157
158
159
160
161
162
class AVLTree(BinarySearchTree[T]):
    """A strictly typed, self-balancing AVL Tree guaranteeing O(log n) operations."""
    def _get_height(self, node: Optional[TreeNode[T]]) -> int:
        if not node:
            return 0
        return node.height

    def _get_balance(self, node: Optional[TreeNode[T]]) -> int:
        if not node:
            return 0
        return self._get_height(node.left) - self._get_height(node.right)

    def _right_rotate(self, y: TreeNode[T]) -> TreeNode[T]:
        """Rotates the subtree to the right."""
        x = y.left
        assert x is not None
        T2 = x.right

        # Perform rotation
        x.right = y
        y.left = T2

        # Update heights
        y.height = 1 + max(self._get_height(y.left), self._get_height(y.right))
        x.height = 1 + max(self._get_height(x.left), self._get_height(x.right))
        return x

    def _left_rotate(self, x: TreeNode[T]) -> TreeNode[T]:
        """Rotates the subtree to the left."""
        y = x.right
        assert y is not None
        T2 = y.left

        # Perform rotation
        y.left = x
        x.right = T2

        # Update heights
        x.height = 1 + max(self._get_height(x.left), self._get_height(x.right))
        y.height = 1 + max(self._get_height(y.left), self._get_height(y.right))
        return y

    def insert(self, value: T) -> None:
        """Public insert method that triggers the recursive balancing engine."""
        self.root = self._insert_recursive(self.root, value)

    def _insert_recursive(self, node: Optional[TreeNode[T]], value: T) -> TreeNode[T]:
        # 1. Normal BST Insert
        if not node:
            return TreeNode(value)
        elif value < node.value:
            node.left = self._insert_recursive(node.left, value)
        elif value > node.value:
            node.right = self._insert_recursive(node.right, value)
        else:
            return node  # Duplicate keys not allowed

        # 2. Update height of current node
        node.height = 1 + max(self._get_height(node.left), self._get_height(node.right))

        # 3. Get the balance factor
        balance = self._get_balance(node)

        # 4. Mathematically balance the tree (The 4 Cases)
        if balance > 1:
            assert node.left is not None
            # Left Left Case
            if value < node.left.value:
                return self._right_rotate(node)
            # Left Right Case
            if value > node.left.value:
                node.left = self._left_rotate(node.left)
                return self._right_rotate(node)

        if balance < -1:
            assert node.right is not None
            # Right Right Case
            if value > node.right.value:
                return self._left_rotate(node)
            # Right Left Case
            if value < node.right.value:
                node.right = self._right_rotate(node.right)
                return self._left_rotate(node)

        return node

    def delete(self, key: T) -> None:
        self.root = self._delete_recursive(self.root, key)

    def _delete_recursive(
        self, root: Optional[TreeNode[T]], key: T
    ) -> Optional[TreeNode[T]]:
        # 1. Standard BST delete
        if not root:
            return root
        if key < root.value:
            root.left = self._delete_recursive(root.left, key)
        elif key > root.value:
            root.right = self._delete_recursive(root.right, key)
        else:
            if not root.left:
                return root.right
            if not root.right:
                return root.left
            temp = self._get_min_value_node(root.right)
            root.value = temp.value
            root.right = self._delete_recursive(root.right, temp.value)

        # 2. Update height and rebalance
        root.height = 1 + max(self._get_height(root.left), self._get_height(root.right))
        balance = self._get_balance(root)

        # Left Left
        if balance > 1:
            assert root.left is not None
            if self._get_balance(root.left) >= 0:
                return self._right_rotate(root)
            # Left Right
            if self._get_balance(root.left) < 0:
                root.left = self._left_rotate(root.left)
                return self._right_rotate(root)

        if balance < -1:
            assert root.right is not None
            # Right Right
            if self._get_balance(root.right) <= 0:
                return self._left_rotate(root)
            # Right Left
            if self._get_balance(root.right) > 0:
                root.right = self._right_rotate(root.right)
                return self._left_rotate(root)

        return root

    def _get_min_value_node(self, node: TreeNode[T]) -> TreeNode[T]:
        if node.left is None:
            return node
        return self._get_min_value_node(node.left)

    def search(self, key: T) -> Optional[T]:
        """Searches for a key and returns it if found."""
        current = self.root
        while current:
            if key == current.value:
                return current.value
            elif key < current.value:
                current = current.left
            else:
                current = current.right
        return None

insert(value)

Public insert method that triggers the recursive balancing engine.

Source code in pure_python_ds/trees/avl_tree.py
55
56
57
def insert(self, value: T) -> None:
    """Public insert method that triggers the recursive balancing engine."""
    self.root = self._insert_recursive(self.root, value)

search(key)

Searches for a key and returns it if found.

Source code in pure_python_ds/trees/avl_tree.py
152
153
154
155
156
157
158
159
160
161
162
def search(self, key: T) -> Optional[T]:
    """Searches for a key and returns it if found."""
    current = self.root
    while current:
        if key == current.value:
            return current.value
        elif key < current.value:
            current = current.left
        else:
            current = current.right
    return None

Red-Black Tree

Bases: Generic[T]

Source code in pure_python_ds/trees/red_black_tree.py
  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
class RedBlackTree(Generic[T]):
    def __init__(self):
        self.NULL = RBNode(None)
        self.NULL.color = Color.BLACK
        self.root = self.NULL
    def __str__(self) -> str:
        """Returns the ASCII visualization of the tree."""
        # Assuming your tree class stores the root node in 'self.root'
        return visualize_binary_tree(self.root)

    def insert(self, key: T):
        node = RBNode(key)
        node.parent = None
        node.value = key
        node.left = self.NULL
        node.right = self.NULL
        node.color = Color.RED

        y = None
        x = self.root

        while x != self.NULL:
            y = x
            if node.value < x.value:
                x = x.left
            else:
                x = x.right

        node.parent = y
        if y is None:
            self.root = node
        elif node.value < y.value:
            y.left = node
        else:
            y.right = node

        if node.parent is None:
            node.color = Color.BLACK
            return

        if node.parent.parent is None:
            return

        self._fix_insert(node)

    def _fix_insert(self, k):
        while k.parent.color == Color.RED:
            if k.parent == k.parent.parent.right:
                u = k.parent.parent.left
                if u.color == Color.RED:
                    u.color = Color.BLACK
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    k = k.parent.parent
                else:
                    if k == k.parent.left:
                        k = k.parent
                        self._right_rotate(k)
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    self._left_rotate(k.parent.parent)
            else:
                u = k.parent.parent.right
                if u.color == Color.RED:
                    u.color = Color.BLACK
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    k = k.parent.parent
                else:
                    if k == k.parent.right:
                        k = k.parent
                        self._left_rotate(k)
                    k.parent.color = Color.BLACK
                    k.parent.parent.color = Color.RED
                    self._right_rotate(k.parent.parent)
            if k == self.root:
                break
        self.root.color = Color.BLACK

    def _left_rotate(self, x):
        y = x.right
        x.right = y.left
        if y.left != self.NULL:
            y.left.parent = x

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

    def _right_rotate(self, x):
        y = x.left
        x.left = y.right
        if y.right != self.NULL:
            y.right.parent = x

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

__str__()

Returns the ASCII visualization of the tree.

Source code in pure_python_ds/trees/red_black_tree.py
14
15
16
17
def __str__(self) -> str:
    """Returns the ASCII visualization of the tree."""
    # Assuming your tree class stores the root node in 'self.root'
    return visualize_binary_tree(self.root)

Trie

Source code in pure_python_ds/trees/trie.py
 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
class Trie:
    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str):
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end_of_word = True

    def search(self, word: str) -> bool:
        node = self.root
        for char in word:
            if char not in node.children:
                return False
            node = node.children[char]
        return node.is_end_of_word

    def starts_with(self, prefix: str) -> bool:
        node = self.root
        for char in prefix:
            if char not in node.children:
                return False
            node = node.children[char]
        return True

    def delete(self, word: str):
        def _delete(node, word, depth):
            if depth == len(word):
                if node.is_end_of_word:
                    node.is_end_of_word = False
                return len(node.children) == 0

            char = word[depth]
            if char not in node.children:
                return False

            should_delete_child = _delete(node.children[char], word, depth + 1)

            if should_delete_child:
                del node.children[char]
                return len(node.children) == 0 and not node.is_end_of_word
            return False

        _delete(self.root, word, 0)