stl_treemap.tree_map

  1"""Ordered map backed by a red-black tree."""
  2
  3from __future__ import annotations
  4
  5from collections import UserDict
  6from collections.abc import Callable, Collection, Iterable
  7from typing import Any
  8
  9from stl_treemap.insertion_result import InsertionResult
 10from stl_treemap.iterators import ReverseIterator, TreeIterator
 11from stl_treemap.policies import KeyValuePolicy
 12from stl_treemap.py_iterators import PyIterator, PyReverseIterator
 13from stl_treemap.tree import Tree
 14from stl_treemap.tree_node import TreeNode
 15
 16
 17class TreeMap[K, V](Collection[K]):
 18    """
 19    Ordered map associating unique keys with values, backed by a red-black tree.
 20
 21    Keys are kept in ascending order at all times. Lookup, insertion, and
 22    deletion are all O(log n).
 23
 24    Example::
 25
 26        m = TreeMap()
 27        m[1] = "a"
 28        m[2] = "b"
 29        v = m[1]  # 'a'
 30        for key, value in m:
 31            print(f"key: {key}, value: {value}")
 32    """
 33
 34    def __init__(self, iterable: Iterable[tuple[K, V]] | None = None, **kwargs) -> None:
 35        """
 36        Create an empty map, or pre-populate it from an iterable of (key, value) pairs.
 37
 38        Args:
 39            iterable: Optional iterable of (key, value) tuples.
 40            kwargs: another way to provide key-value pairs
 41
 42        Raises:
 43            TypeError: When *iterable* is not iterable.
 44
 45        Example::
 46
 47            m = TreeMap()
 48            m = TreeMap({2: "B", 1: "A", 3: "C"})
 49            m = TreeMap({2: "B", 1: "A", 3: "C"}.items())
 50            m = TreeMap([[2, "B"], [1, "A"], [3, "C"]])
 51            m = TreeMap(A=1, B=2, C=3)
 52
 53        """
 54        self._t: Tree[K, V] = Tree()
 55        self._t.value_policy = KeyValuePolicy()
 56        if iterable is not None:
 57            self.update(iterable)
 58        if kwargs:
 59            self.update(kwargs.items())
 60
 61    def update(self, iterable: Iterable[tuple[K, V]]) -> None:
 62        """
 63        Add all contents from the iterable of (key, value) pairs.
 64
 65        Args:
 66            iterable: Optional iterable of (key, value) tuples.
 67
 68        Raises:
 69            TypeError: When *iterable* is not iterable.
 70
 71        Example::
 72
 73            TreeMap().update(TreeMap({2: "B", 1: "A", 3: "C"}))
 74            TreeMap().update({2: "B", 1: "A", 3: "C"})
 75            TreeMap().update({2: "B", 1: "A", 3: "C"}.items())
 76            TreeMap().update([[2, "B"], [1, "A"], [3, "C"]])
 77
 78        """
 79        if isinstance(iterable, (dict, UserDict, TreeMap)):
 80            for k, v in iterable.items():
 81                self.set(k, v)
 82        elif hasattr(iterable, "__iter__"):
 83            for k, v in iterable:
 84                self.set(k, v)
 85        else:
 86            raise TypeError("TreeMap accepts only iterable objects")
 87
 88    # ------------------------------------------------------------------
 89    # Python dict-compatible methods
 90    # ------------------------------------------------------------------
 91
 92    def clear(self) -> None:
 93        """
 94        Remove all key-value pairs.
 95
 96        Example::
 97
 98            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
 99            m.clear()
100            len(m)  # 0
101
102        """
103        self._t.clear()
104
105    def delete(self, key: K) -> None:
106        """
107        @private Remove the entry with *key*; does nothing if the key is absent.
108
109        Example::
110
111            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
112            m.delete(2)  # m is now {1:A,3:C}
113            m.delete(99)  # no-op
114
115        """
116        it = self._t.find(key)
117        if not it.equals(self._t.end()):
118            self._t.erase(it.node)
119
120    def get(self, key: K, default: V = None) -> V | None:
121        """
122        Return the value for *key*, or *default* if the key is absent.
123
124        Args:
125            key: Key to look up.
126            default: Value returned when *key* is not in the map.
127
128        Example::
129
130            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
131            m.get(1)  # "A"
132            m.get(4)  # None
133            m.get(4, "Z")  # "Z"
134
135        """
136        it = self._t.find(key)
137        if not it.equals(self._t.end()):
138            return it.value
139        return default
140
141    def has(self, key: K) -> bool:
142        """
143        @private Return True when *key* exists in the map.
144
145        Example::
146
147            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
148            m.has(1)  # True
149            m.has(4)  # False
150
151        """
152        it = self._t.find(key)
153        return not it.equals(self._t.end())
154
155    def set(self, key: K, value: V) -> None:
156        """
157        @private Insert *key*/*value*, or update the value if *key* already exists.
158
159        Example::
160
161            m = TreeMap()
162            m.set(1, "A")
163            m.set(1, "B")  # updates existing entry; m is now {1:B}
164
165        """
166        n: TreeNode[K, V] = TreeNode()
167        n.key = key
168        n.value = value
169        self._t.insert_or_replace(n)
170
171    def items(self) -> PyIterator[tuple[K, V]]:
172        """
173        Return a forward iterator over (key, value) pairs in ascending key order.
174
175        Example:
176            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
177            for key, value in m.items():
178                print(f"key: {key}, value: {value}")
179            # key: 1, value: A
180            # key: 2, value: B
181            # key: 3, value: C
182
183        """
184        return self._t.items()
185
186    def keys(self) -> PyIterator[K]:
187        """
188        Return a forward iterator over keys in ascending order.
189
190        Example::
191
192            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
193            list(m.keys())  # [1, 2, 3]
194            list(m.keys().backwards())  # [3, 2, 1]
195
196        """
197        return self._t.keys()
198
199    def values(self) -> PyIterator[V]:
200        """
201        Return a forward iterator over values in ascending key order.
202
203        Example::
204
205            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
206            list(m.values())  # ["A", "B", "C"]
207            list(m.values().backwards())  # ["C", "B", "A"]
208
209        """
210        return self._t.values()
211
212    @property
213    def size(self) -> int:
214        """
215        @private Number of entries in the map.
216
217        Example::
218
219            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
220            len(m)  # 3
221
222        """
223        return self._t.size()
224
225    def __getitem__(self, key: K) -> V | None:
226        """
227        @public Return the value associated with key.
228
229        Args:
230            key: Key to look up.
231
232        Returns:
233            Value stored at key.
234
235        Raises:
236            KeyError: When key is not present in the tree.
237
238        Example::
239
240            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
241            m[1]  # "A"
242            m[4]  # raises KeyError
243
244        """
245        it = self._t.find(key)
246        if not it.equals(self._t.end()):
247            return it.value
248        raise KeyError(f"Key {key} not found")
249
250    def __setitem__(self, key: K, value: V) -> None:
251        """
252        @public Insert or replace the value for key.
253
254        Args:
255            key: Key to insert or update.
256            value: Value to associate with the key.
257
258        Example::
259
260            m = TreeMap()
261            m[1] = "A"
262            m[1] = "B"  # updates existing entry
263
264        """
265        self.set(key, value)
266
267    def __contains__(self, key: K) -> bool:
268        """
269        @public Return true if tree contains provided key.
270
271        Args:
272            key: Key to check.
273
274        Example::
275
276            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
277            1 in m  # True
278            4 in m  # False
279
280        """
281        iter = self.find(key)
282        return not iter.equals(self.end())
283
284    def __iter__(self) -> PyIterator[K]:
285        """
286        @public Iterate over keys in ascending key order.
287
288        Example::
289
290            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
291            for key in m:
292                print(f"{key}")
293
294        """
295        return self._t.keys()
296
297    def __len__(self) -> int:
298        """
299        @public Return the number of entries in the map.
300
301        Example::
302
303            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
304            len(m)  # 3
305
306        """
307        return self._t.size()
308
309    def __str__(self) -> str:
310        """
311        @public Return a string representation in the form {key1:value1,key2:value2,...}.
312
313        Example::
314
315            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
316            str(m)  # "{1:A,2:B,3:C}"
317
318        """
319        return str(self._t)
320
321    def __repr__(self) -> str:
322        """
323        @public Return a string representation of the container's contents.
324
325        Example::
326
327            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
328            repr(m)  # "{1:A,2:B,3:C}"
329
330        """
331        return self.__str__()
332
333    def __delitem__(self, key: K):
334        """
335        @public Delete item by key or raise KeyError.
336
337        Args:
338            key: Key to remove.
339
340        Raises:
341            KeyError: When key is not present in the map.
342
343        Example::
344
345            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
346            del m[2]  # m is now {1:A,3:C}
347            del m[99]  # raises KeyError
348
349        """
350        del self._t[key]
351
352    def popitem(self, key: K) -> tuple[K, V]:
353        """
354        Remove *key* and return its key-value tuple or raise KeyError.
355
356        Args:
357            key: Key to remove.
358
359        Example::
360
361            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
362            m.popitem(2)  # "(2,B)", m is now {1:A,3:C}
363            m.popitem(9)  # raises KeyError
364
365        """
366        iter = self.find(key)
367        if not iter.equals(self.end()):
368            k = iter.key
369            v = iter.value
370            self.erase(iter)
371            return (k, v)
372        raise KeyError(f"Key {key} not found")
373
374    def pop(self, key: K, default: K | None = None) -> V:
375        """
376        Remove *key* and return its value, or return *default* if absent.
377
378        Args:
379            key: Key to remove.
380            default: Value to return when *key* is not present.
381
382        Example::
383
384            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
385            m.pop(2)  # "B", m is now {1:A,3:C}
386            m.pop(9)  # None, m is unchanged
387            m.pop(9, "Z")  # "Z", m is unchanged
388
389        """
390        iter = self.find(key)
391        if not iter.equals(self.end()):
392            res = iter.value
393            self.erase(iter)
394            return res
395        return default
396
397    # ------------------------------------------------------------------
398    # Additional iteration
399    # ------------------------------------------------------------------
400
401    def backwards(self) -> PyReverseIterator[K]:
402        """
403        Return a reverse iterator over keys in descending key order.
404
405        Example::
406
407            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
408            for key in m.backwards():
409                print(f"{key}")
410            # 3
411            # 2
412            # 1
413
414        """
415        return self._t.keys().backwards()
416
417    # ------------------------------------------------------------------
418    # Custom comparator
419    # ------------------------------------------------------------------
420
421    @property
422    def compare_func(self) -> Callable[[Any, Any], int]:
423        """
424        @private The current 3-way comparison function is used to order keys.
425
426        The function must accept two key arguments and return a negative
427        integer, zero, or a positive integer when the first key is less
428        than, equal to, or greater than the second key.
429
430        Example::
431
432            m = TreeMap([[3, "C"], [1, "A"], [2, "B"]])
433            m.compare_func = lambda a, b: -1 if a < b else (1 if a > b else 0)
434            m.set(4, "D")
435            list(m.keys())  # [1, 2, 3, 4]
436
437        """
438        return self._t.compare
439
440    @compare_func.setter
441    def compare_func(self, func: Callable[[Any, Any], int]) -> None:
442        """@private Replace the comparison function. Clears the map because existing order is invalid."""
443        self.clear()
444        self._t.compare = func
445
446    # ------------------------------------------------------------------
447    # STL-style iterators
448    # ------------------------------------------------------------------
449
450    def begin(self) -> TreeIterator[K, V]:
451        """
452        Return a forward STL-like iterator to the first entry.
453
454        Example::
455
456            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
457            it = m.begin()
458            while not it.equals(m.end()):
459                print(f"key: {it.key}, value: {it.value}")
460                it.next()
461
462        """
463        return self._t.begin()
464
465    def end(self) -> TreeIterator[K, V]:
466        """
467        Return a forward STL-like iterator past the last entry.
468
469        Example::
470
471            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
472            it = m.begin()
473            while not it.equals(m.end()):
474                print(f"key: {it.key}, value: {it.value}")
475                it.next()
476
477        """
478        return self._t.end()
479
480    def find(self, key: K) -> TreeIterator[K, V]:
481        """
482        Return an STL-like iterator to the entry with *key*, or end() if not found.
483
484        Example::
485
486            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
487            it = m.find(2)
488            if not it.equals(m.end()):
489                print(f"key: {it.key}, value: {it.value}")  # key: 2, value: B
490            it = m.find(99)
491            it.equals(m.end())  # True
492
493        """
494        return self._t.find(key)
495
496    def insert_unique(self, key: K, value: V) -> InsertionResult[TreeIterator[K, V]]:
497        """
498        Insert *(key, value)* only when *key* is not already present.
499
500        Returns:
501            InsertionResult with was_added=True and an iterator when the key
502            was new; was_added=False and iterator=None when the key existed.
503
504        Example::
505
506            m = TreeMap()
507            res = m.insert_unique(1, "A")
508            res.was_added  # True
509            res.iterator.value  # "A"
510            res2 = m.insert_unique(1, "B")
511            res2.was_added  # False — key already exists, no change made
512
513        """
514        n: TreeNode[K, V] = TreeNode()
515        n.key = key
516        n.value = value
517        return self._t.insert_unique(n)
518
519    def insert_or_replace(self, key: K, value: V) -> InsertionResult[TreeIterator[K, V]]:
520        """
521        Insert *(key, value)*, replacing the existing value when *key* already exists.
522
523        Returns:
524            InsertionResult with was_added=True on new insertion, or
525            was_replaced=True and the updated iterator on replacement.
526
527        Example::
528
529            m = TreeMap()
530            res = m.insert_or_replace(1, "A")
531            res.was_added  # True
532            res2 = m.insert_or_replace(1, "B")
533            res2.was_replaced  # True
534            res2.iterator.value  # "B"
535
536        """
537        n: TreeNode[K, V] = TreeNode()
538        n.key = key
539        n.value = value
540        return self._t.insert_or_replace(n)
541
542    def erase(self, iterator: TreeIterator[K, V]) -> None:
543        """
544        Remove the entry pointed to by *iterator*.
545
546        Example::
547
548            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
549            it = m.find(2)
550            m.erase(it)
551            str(m)  # "{1:A,3:C}"
552
553        """
554        self._t.erase(iterator.node)
555
556    def lower_bound(self, key: K) -> TreeIterator[K, V]:
557        """
558        Return an STL-like iterator to the first entry with key >= *key*, or end().
559
560        Example::
561
562            m = TreeMap([[2, "B"], [4, "D"], [6, "F"], [8, "H"]])
563            lo = m.lower_bound(3)  # iterator to key 4
564            hi = m.upper_bound(6)  # iterator to key 8
565            it = lo
566            while not it.equals(hi):
567                print(it.key)  # 4, 6
568                it.next()
569
570        """
571        return self._t.lower_bound(key)
572
573    def rbegin(self) -> ReverseIterator[K, V]:
574        """
575        Return a reverse STL-like iterator to the entry with largest key.
576
577        Example::
578
579            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
580            it = m.rbegin()
581            while not it.equals(m.rend()):
582                print(f"key: {it.key}, value: {it.value}")
583                it.next()
584            # key: 3, value: C
585            # key: 2, value: B
586            # key: 1, value: A
587
588        """
589        return self._t.rbegin()
590
591    def rend(self) -> ReverseIterator[K, V]:
592        """
593        Return a reverse STL-like iterator before the entry with the smallest key.
594
595        Example::
596
597            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
598            it = m.rbegin()
599            while not it.equals(m.rend()):
600                print(f"key: {it.key}, value: {it.value}")
601                it.next()
602
603        """
604        return self._t.rend()
605
606    def upper_bound(self, key: K) -> TreeIterator[K, V]:
607        """
608        Return an STL-like iterator to the first entry with key larger than *key*, or end().
609
610        Example::
611
612            m = TreeMap([[2, "B"], [4, "D"], [6, "F"], [8, "H"]])
613            lo = m.lower_bound(3)  # iterator to key 4
614            hi = m.upper_bound(6)  # iterator to key 8
615            it = lo
616            while not it.equals(hi):
617                print(it.key)  # 4, 6
618                it.next()
619
620        """
621        return self._t.upper_bound(key)
622
623    def first(self) -> tuple[K, V] | None:
624        """
625        Return the first (key, value) pair, or None when the map is empty.
626
627        Example::
628
629            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
630            key, value = m.first()  # (1, "A")
631            TreeMap().first()  # None
632
633        """
634        return self._t.first()
635
636    def last(self) -> tuple[K, V] | None:
637        """
638        Return the last (key, value) pair, or None when the map is empty.
639
640        Example::
641
642            m = TreeMap([[1, "A"], [2, "B"], [3, "C"]])
643            key, value = m.last()  # (3, "C")
644            TreeMap().last()  # None
645
646        """
647        return self._t.last()
648
649    def __or__(self, other: Any) -> TreeMap[K, V]:
650        """
651        @public Return a new map that is the union of this map and *other*.
652
653        Values from *other* override values in this map for shared keys.
654
655        Example::
656
657            m1 = TreeMap([[1, "A"], [2, "B"]])
658            m2 = TreeMap([[2, "X"], [3, "C"]])
659            m3 = m1 | m2
660            str(m3)  # "{1:A,2:X,3:C}"
661
662        """
663        res = TreeMap[K, V](self.items())
664        if isinstance(other, (UserDict, dict, TreeMap)) or hasattr(other, "__iter__"):
665            res.update(other)
666        else:
667            return NotImplemented
668        return res
669
670    def __ror__(self, other):
671        """
672        @public Return a union map when this TreeMap is on the right-hand side of ``|``.
673
674        Example::
675
676            m = TreeMap([[2, "X"], [3, "C"]])
677            result = {1: "A", 2: "B"} | m
678            str(result)  # "{1:A,2:X,3:C}"
679
680        """
681        res = TreeMap[K, V]()
682        if isinstance(other, (UserDict, dict, TreeMap)) or hasattr(other, "__iter__"):
683            res.update(other)
684        else:
685            return NotImplemented
686        res.update(self)
687        return res
688
689    def __ior__(self, other: Any) -> TreeMap[K, V]:
690        """
691        @public Update this map in-place with key-value pairs from *other*.
692
693        Example::
694
695            m = TreeMap([[1, "A"], [2, "B"]])
696            m |= {2: "X", 3: "C"}
697            str(m)  # "{1:A,2:X,3:C}"
698
699        """
700        if isinstance(other, (UserDict, dict, TreeMap)) or hasattr(other, "__iter__"):
701            self.update(other)
702        else:
703            return NotImplemented
704        return self
705
706    def __copy__(self):
707        """
708        @public Create a shallow copy of the map.
709
710        Example::
711
712            import copy
713
714            m = TreeMap([[1, "A"], [2, "B"]])
715            m2 = copy.copy(m)
716            m2[3] = "C"
717            str(m)  # "{1:A,2:B}"      — original unchanged
718            str(m2)  # "{1:A,2:B,3:C}"
719
720        """
721        res = TreeMap[K, V](self.items())
722        return res