15class AvlTree(MutableMapping[_K, _V]):
 16    """Lightweight, pure-python AVL tree.
 17
 18    This class implements the MutableMapping interface and can be used in almost every
 19    way that a built-in dictionary can.
 20
 21    # Sorting
 22    The AVL tree data structure makes it possible to easily iterate on keys in
 23    sort-order, and any method which returns an iterable of keys, values, or items will
 24    return them in sort-order by key.
 25
 26    # Keys
 27    Anything used as a key in an AvlTree must implement `__eq__`, `__hash__`, and
 28    `__lt__`. That is to say they must be immutable and have a less-than comparison.
 29
 30    It's recommended to only insert keys which are all the same type, ensuring that they
 31    have a well-ordering. However, keys of different types can be inserted as long as
 32    their `__eq__` and `__lt__` methods behave.
 33
 34    This class provides no protections against poorly behaved keys and can fail in an
 35    undefined manner if keys are not implemented properly.
 36
 37    # Values
 38    Values can be any Python object.
 39
 40    # Recursion
 41    This class does not use recursive techniques. This ensures that this package can be
 42    used on platforms with low recursion limits, even in scenarios with very large and
 43    very deep trees.
 44
 45    # Time Complexities
 46    Time complexities for specific methods can be found in their docstrings.
 47    """
 48
 49    def __init__(self, mapping: Mapping[_K, _V] | None = None) -> None:
 50        """Constructor.
 51
 52        Inserting the elements of a passed-in mapping has an amortized and worst-case
 53        time complexity of O[n*log(n)].
 54
 55        Example usage:
 56
 57        >>> from avltree import AvlTree
 58        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
 59        >>> avl_tree
 60        AvlTree({0: 'a', 1: 'b'})
 61
 62        Args:
 63            mapping (dict[_K, _V] | None): An optional initial mapping of items to add
 64                to this tree. Defaults to None.
 65        """
 66        self.__nodes: Final[dict[_K, AvlTreeNode[_K, _V]]] = {}
 67        self.__root_key: _K | None = None
 68        if mapping is not None:
 69            for key, value in mapping.items():
 70                self[key] = value
 71
 72    def __setitem__(self, __k: _K, __v: _V) -> None:
 73        """Maps the given key to the given value in this tree.
 74
 75        This method has an amortized and worst-case time complexity of O[log(n)].
 76
 77        Example usage:
 78
 79        >>> from avltree import AvlTree
 80        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
 81        >>> avl_tree[0] = 'foo'
 82        >>> avl_tree[2] = 'c'
 83        >>> avl_tree
 84        AvlTree({0: 'foo', 1: 'b', 2: 'c'})
 85
 86        Args:
 87            __k (_K): The key to map.
 88            __v (_V): The value to associate with the key.
 89        """
 90        if __k in self.__nodes:
 91            self.__nodes[__k].value = __v
 92            return
 93        if self.__root_key is None:
 94            self.__root_key = __k
 95            self.__nodes[self.__root_key] = AvlTreeNode[_K, _V](value=__v)
 96            return
 97        stack: Final[list[_K]] = [self.__root_key]
 98        current_node: AvlTreeNode[_K, _V] = self.__nodes[stack[-1]]
 99        while True:
100            if __k < stack[-1] and current_node.lesser_child_key is None:
101                current_node.lesser_child_key = __k
102                self.__nodes[__k] = AvlTreeNode[_K, _V](value=__v)
103                break
104            if stack[-1] < __k and current_node.greater_child_key is None:
105                current_node.greater_child_key = __k
106                self.__nodes[__k] = AvlTreeNode[_K, _V](value=__v)
107                break
108            if __k < stack[-1] and current_node.lesser_child_key is not None:
109                stack.append(current_node.lesser_child_key)
110                current_node = self.__nodes[stack[-1]]
111            elif current_node.greater_child_key is not None:
112                stack.append(current_node.greater_child_key)
113                current_node = self.__nodes[stack[-1]]
114        self.__enforce_avl(stack=stack)
115
116    def __delitem__(self, __k: _K) -> None:  # noqa: C901, PLR0912
117        """Deletes the given key from this tree.
118
119        This method has an amortized and worst-case time complexity of O[log(n)].
120
121        Example usage:
122
123        >>> from avltree import AvlTree
124        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
125        >>> del avl_tree[0]
126        >>> avl_tree
127        AvlTree({1: 'b'})
128
129        Args:
130            __k (_K): The key to delete from this tree.
131
132        Raises:
133            KeyError: If the given key is not present in this tree.
134        """
135        if self.__root_key is None:
136            message: str = f"Key not present in AvlTree: {__k!r}"
137            raise KeyError(message)
138
139        # Find the node to discard and its parent node
140        parent: AvlTreeNode[_K, _V] | None = None
141        node_key: _K = self.__root_key
142        stack: Final[list[_K]] = [node_key]
143        node: AvlTreeNode[_K, _V] = self.__nodes[node_key]
144        while node_key != __k:
145            parent = node
146            node_key = self.__get_closer_key(key=__k, current_key=node_key)
147            stack.append(node_key)
148            node = self.__nodes[node_key]
149
150        # Find the key of the node with which to replace the existing node
151        replacement_key: _K | None = None
152        if node.lesser_child_key is not None and node.greater_child_key is None:
153            replacement_key = node.lesser_child_key
154        elif node.lesser_child_key is None and node.greater_child_key is not None:
155            replacement_key = node.greater_child_key
156        elif node.lesser_child_key is not None and node.greater_child_key is not None:
157            # Find the next highest node than the one to remove
158            successor_parent: AvlTreeNode[_K, _V] | None = None
159            replacement_key = node.greater_child_key
160            successor: AvlTreeNode[_K, _V] = self.__nodes[replacement_key]
161            while successor.lesser_child_key is not None:
162                successor_parent = successor
163                stack.append(replacement_key)
164                replacement_key = successor.lesser_child_key
165                successor = self.__nodes[successor.lesser_child_key]
166
167            # Swap the successor node with the node to replace
168            if successor_parent is not None and successor.greater_child_key is None:
169                successor_parent.lesser_child_key = None
170                successor.greater_child_key = node.greater_child_key
171            elif successor_parent is not None:
172                successor_parent.lesser_child_key = successor.greater_child_key
173                successor.greater_child_key = node.greater_child_key
174            successor.lesser_child_key = node.lesser_child_key
175
176        # Swap the node to its replacement
177        if parent is None:
178            self.__root_key = replacement_key
179        elif parent.lesser_child_key == node_key:
180            parent.lesser_child_key = replacement_key
181        else:
182            parent.greater_child_key = replacement_key
183        del self.__nodes[node_key]
184        if replacement_key is None:
185            stack.remove(node_key)
186        else:
187            stack[stack.index(node_key)] = replacement_key
188        self.__enforce_avl(stack=stack)
189
190    def __getitem__(self, __k: _K) -> _V:
191        """Gets the value associated with the given key.
192
193        This method has an amortized and worst-case time complexity of O[log(n)].
194
195        Example usage:
196
197        >>> from avltree import AvlTree
198        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
199        >>> avl_tree[0]
200        'a'
201        >>> avl_tree[2]
202        KeyError: 'Key not present in AvlTree: 2'
203
204        Args:
205            __k (_K): The key.
206
207        Returns:
208            _V: The value associated with the given key.
209
210        Raises:
211            KeyError: If the given key is not present in this tree.
212        """
213        if self.__root_key is None:
214            message: str = f"Key not present in AvlTree: {__k!r}"
215            raise KeyError(message)
216        current_key: _K = self.__root_key
217        current_node: AvlTreeNode[_K, _V] = self.__nodes[current_key]
218        while current_key != __k:
219            current_key = self.__get_closer_key(key=__k, current_key=current_key)
220            current_node = self.__nodes[current_key]
221        return current_node.value
222
223    def __len__(self) -> int:
224        """Returns the number of items contained in this tree.
225
226        This method has an amortized and worst-case time complexity of O[1].
227
228        Example usage:
229
230        >>> from avltree import AvlTree
231        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
232        >>> len(avl_tree)
233        2
234
235        Returns:
236            int: The number of items contained in this tree.
237        """
238        return len(self.__nodes)
239
240    def __iter__(self) -> Iterator[_K]:
241        """Iterates over all keys contained in this tree in sort order.
242
243        Getting the first key has an amortized and worst-case time complexity of
244        O[log(n)]. Iterating over all keys has an amortized and worst-case time
245        complexity of O[n].
246
247        Example usage:
248
249        >>> from avltree import AvlTree
250        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
251        >>> for key in avl_tree:
252        ...     print(key)
253        ...
254        0
255        1
256
257        Returns:
258            Iterator[_K]: The iterator object.
259        """
260        return self.between()
261
262    def __repr__(self) -> str:
263        """Builds a developer-friendly string representation of this AvlTree.
264
265        Example usage:
266
267        >>> from avltree import AvlTree
268        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
269        >>> repr(avl_tree)
270        "AvlTree({0: 'a', 1: 'b'})"
271
272
273        Returns:
274            str: A string representation of this AvlTree.
275        """
276        return f"AvlTree({{{', '.join(f'{k!r}: {v!r}' for k, v in self.items())}}})"
277
278    def minimum(self) -> _K:
279        """Gets the minimum key contained in this tree.
280
281        This method has an amortized and worst-case time complexity of O[log(n)].
282
283        Example usage:
284
285        >>> from avltree import AvlTree
286        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
287        >>> avl_tree.minimum()
288        0
289
290        Returns:
291            _K: The minimum key.
292
293        Raises:
294            ValueError: If there are no keys present in this tree.
295        """
296        if self.__root_key is None:
297            message: Final[str] = "Cannot get the minimum of an empty AvlTree"
298            raise ValueError(message)
299        key: _K = self.__root_key
300        while self.__nodes[key].lesser_child_key is not None:
301            key = self.__nodes[key].lesser_child_key  # type: ignore[assignment]
302        return key
303
304    def maximum(self) -> _K:
305        """Gets the maximum key contained in this tree.
306
307        This method has an amortized and worst-case time complexity of O[log(n)].
308
309        Example usage:
310
311        >>> from avltree import AvlTree
312        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b'})
313        >>> avl_tree.maximum()
314        1
315
316        Returns:
317            _K: The maximum key.
318
319        Raises:
320            ValueError: If there are no keys present in this tree.
321        """
322        if self.__root_key is None:
323            message: Final[str] = "Cannot get the maximum of an empty AvlTree"
324            raise ValueError(message)
325        key: _K = self.__root_key
326        node: AvlTreeNode[_K, _V] = self.__nodes[key]
327        while node.greater_child_key is not None:
328            key = node.greater_child_key
329            node = self.__nodes[key]
330        return key
331
332    def between(  # noqa: C901, PLR0912
333        self,
334        start: _K | None = None,
335        stop: _K | None = None,
336        treatment: Literal["inclusive", "exclusive"] = "inclusive",
337    ) -> Iterator[_K]:
338        """Iterates over all keys between the given start and stop in sort order.
339
340        Getting the first key has an amortized and worst-case time complexity of
341        O[log(n)]. Iterating over all keys has an amortized and worst-case time
342        complexity of O[k], where k is the number of items in only the interval between
343        start and stop.
344
345        Example usage:
346
347        >>> from avltree import AvlTree
348        >>> avl_tree = AvlTree[int, str]({0: 'a', 1: 'b', 2: 'c'})
349        >>> for key in avl_tree.between(start=0, stop=2, treatment="inclusive"):
350        ...     print(key)
351        ...
352        0
353        1
354        2
355        >>> for key in avl_tree.between(start=0, treatment="exclusive"):
356        ...     print(key)
357        ...
358        1
359        2
360
361        Args:
362            start (_K | None): The key at which to start iterating. If None, iteration
363                starts at the minimum key. Defaults to None.
364            stop (_K | None): The key at which to stop iterating. If None, iteration
365                stops at the maximum key. Defaults to None.
366            treatment (Literal["inclusive", "exclusive"]): Whether the given start and
367                stop should be included or excluded from the returned iterator. Has no
368                effect when start and stop are None. Defaults to "inclusive".
369
370        Returns:
371            Iterator[_K]: An iterator which will iterate over all keys between the given
372                start and stop.
373        """
374        if self.__root_key is None:
375            return
376        stack: Final[list[tuple[_K, bool]]] = []
377        current_key: _K = self.__root_key
378        while True:
379            current_node: AvlTreeNode[_K, _V] = self.__nodes[current_key]
380            if start is None or start < current_key:
381                stack.append((current_key, True))
382                if current_node.lesser_child_key is not None:
383                    current_key = current_node.lesser_child_key
384                else:
385                    break
386            elif start == current_key:
387                if treatment == "inclusive":
388                    stack.append((current_key, True))
389                    break
390                elif current_node.greater_child_key is not None:
391                    current_key = current_node.greater_child_key
392                else:
393                    break
394            elif current_node.greater_child_key is not None:
395                current_key = current_node.greater_child_key
396            else:
397                break
398        while len(stack) > 0:
399            key, lesser_child_visited = stack.pop(-1)
400            node: AvlTreeNode[_K, _V] = self.__nodes[key]
401            if (
402                stop is not None  # noqa: PLR0916
403                and (stop < key or (stop == key and treatment == "exclusive"))
404                and (lesser_child_visited or node.lesser_child_key is None)
405            ):
406                break
407            elif node.lesser_child_key is not None and not lesser_child_visited:
408                stack.append((key, True))
409                stack.append((node.lesser_child_key, False))
410            elif node.greater_child_key is not None:
411                stack.append((node.greater_child_key, False))
412                yield key
413            else:
414                yield key
415
416    def __get_closer_key(self, key: _K, current_key: _K) -> _K:
417        """Gets the next closest key to the given key.
418
419        Args:
420            key (_K): The key to search for.
421            current_key (_K): The current key.
422
423        Returns:
424            _K: The next closest key to the given key.
425
426        Raises:
427            KeyError: If the given key is not present in this tree.
428        """
429        current_node: Final[AvlTreeNode[_K, _V]] = self.__nodes[current_key]
430        if key < current_key and current_node.lesser_child_key is not None:
431            return current_node.lesser_child_key
432        if current_key < key and current_node.greater_child_key is not None:
433            return current_node.greater_child_key
434        message: Final[str] = f"Key not present in AvlTree: {key!r}"
435        raise KeyError(message)
436
437    def __enforce_avl(self, stack: list[_K]) -> None:
438        """Enforces the AVL property on this tree.
439
440        Args:
441            stack (list[_K]): The stack to traverse in reverse order.
442        """
443        while len(stack) > 0:
444            key: _K = stack.pop(-1)
445            node: AvlTreeNode[_K, _V] = self.__nodes[key]
446            balance: int = self.__calculate_balance(node=node)
447            if -1 <= balance <= 1:
448                self.__update_height(node=node)
449                continue
450            if balance == -2:  # noqa: PLR2004
451                lesser_child_key: _K = cast(_K, node.lesser_child_key)
452                if self.__calculate_balance(node=self.__nodes[lesser_child_key]) == 1:
453                    node.lesser_child_key = self.__rotate(
454                        key=lesser_child_key,
455                        direction="left",
456                    )
457                replacement_key: _K = self.__rotate(key=key, direction="right")
458            else:
459                greater_child_key: _K = cast(_K, node.greater_child_key)
460                if self.__calculate_balance(node=self.__nodes[greater_child_key]) == -1:
461                    node.greater_child_key = self.__rotate(
462                        key=greater_child_key,
463                        direction="right",
464                    )
465                replacement_key = self.__rotate(key=key, direction="left")
466            parent_node: AvlTreeNode[_K, _V] | None = (
467                None if len(stack) == 0 else self.__nodes[stack[-1]]
468            )
469            if parent_node is None:
470                self.__root_key = replacement_key
471            elif parent_node.lesser_child_key == key:
472                parent_node.lesser_child_key = replacement_key
473            else:
474                parent_node.greater_child_key = replacement_key
475
476    def __calculate_balance(self, node: AvlTreeNode[_K, _V]) -> int:
477        """Calculates the balance factor of the given node.
478
479        Args:
480            node (AvlTreeNode[_K, _V]): The node.
481
482        Returns:
483            int: The balance factor of the given node.
484        """
485        return (
486            -1
487            if node.greater_child_key is None
488            else self.__nodes[node.greater_child_key].height
489        ) - (
490            -1
491            if node.lesser_child_key is None
492            else self.__nodes[node.lesser_child_key].height
493        )
494
495    def __update_height(self, node: AvlTreeNode[_K, _V]) -> None:
496        """Updates the height of the given node.
497
498        Args:
499            node (AvlTreeNode[_K, _V]): The node.
500        """
501        node.height = 1 + max(
502            (
503                -1
504                if node.greater_child_key is None
505                else self.__nodes[node.greater_child_key].height
506            ),
507            (
508                -1
509                if node.lesser_child_key is None
510                else self.__nodes[node.lesser_child_key].height
511            ),
512        )
513
514    def __rotate(self, key: _K, direction: Literal["left", "right"]) -> _K:
515        """Performs a rotation at the given key.
516
517        Args:
518            key (_K): The key to perform a right rotation on.
519            direction (Literal["left", "right"]): The direction of the rotation.
520
521        Returns:
522            _K: The new root key of this subtree.
523
524        Raises:
525            ValueError: If the shape of the tree is incompatible with the requested
526                rotation direction.
527        """
528        node: Final[AvlTreeNode[_K, _V]] = self.__nodes[key]
529        replacement_key: Final[_K] = cast(
530            _K,
531            node.greater_child_key if direction == "left" else node.lesser_child_key,
532        )
533        replacement_node: Final[AvlTreeNode[_K, _V]] = self.__nodes[replacement_key]
534        if direction == "left":
535            node.greater_child_key = replacement_node.lesser_child_key
536            replacement_node.lesser_child_key = key
537        else:
538            node.lesser_child_key = replacement_node.greater_child_key
539            replacement_node.greater_child_key = key
540        self.__update_height(node=node)
541        self.__update_height(node=replacement_node)
542        return replacement_key