stl_treemap.tree_multimap

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