stl_treemap.tree_multiset

  1"""Ordered multiset backed by a red-black tree (duplicate keys allowed)."""
  2
  3from __future__ import annotations
  4
  5from collections.abc import Callable, Iterable
  6from typing import Any
  7
  8from stl_treemap.insertion_result import InsertionResult
  9from stl_treemap.iterators import ReverseIterator, TreeIterator
 10from stl_treemap.policies import KeyOnlyPolicy
 11from stl_treemap.py_iterators import PyIterator, PyReverseIterator
 12from stl_treemap.tree import Tree
 13from stl_treemap.tree_node import TreeNode
 14
 15
 16class TreeMultiSet[K]:
 17    """
 18    Ordered multiset of keys backed by a red-black tree.
 19
 20    Unlike TreeSet, multiple equal keys are allowed. Elements are kept in
 21    ascending order. Lookup, insertion, and deletion are all O(log n).
 22
 23    Example::
 24
 25        s = TreeMultiSet()
 26        s.add(1)
 27        s.add(2)
 28        s.add(2)
 29        for key in s:
 30            print(key)  # 1, 2, 2
 31    """
 32
 33    def __init__(self, iterable: Iterable[K] | None = None, *args) -> None:
 34        """
 35        Create an empty multiset, or pre-populate it from an iterable of keys.
 36
 37        Duplicate keys are preserved (multiset semantics).
 38
 39        Args:
 40            iterable: Optional iterable of keys.
 41            args: initial keys could be provided as extra parameters
 42
 43        Raises:
 44            TypeError: When *iterable* is not iterable.
 45
 46        Example::
 47
 48            s = TreeMultiSet()
 49            s = TreeMultiSet([1, 2, 2, 3])  # {1,2,2,3}
 50            s = TreeMultiSet({1, 2})  # {1,2}
 51            s = TreeMultiSet(None, 1, 2, 1, 2)  # {1,1,2,2}
 52
 53        """
 54        self._t: Tree[K, K] = Tree()
 55        self._t.value_policy = KeyOnlyPolicy()
 56        if iterable is not None:
 57            self.update(iterable)
 58        if args:
 59            self.update(args)
 60
 61    def update(self, *iterables: Iterable[K]) -> None:
 62        """
 63        Add contents from all provided iterables to the current set.
 64
 65        Duplicate keys are preserved (multiset semantics).
 66
 67        Args:
 68            iterables: One or more iterables containing keys.
 69
 70        Raises:
 71            TypeError: When *iterable* is not iterable.
 72
 73        Example::
 74
 75            s = TreeMultiSet()
 76            s.update([1, 2, 2, 3])  # {1,2,2,3}
 77
 78        """
 79        for iterable in iterables:
 80            if not hasattr(iterable, "__iter__"):
 81                raise TypeError("TreeMultiSet accepts only iterable objects")
 82            for k in iterable:
 83                self.add(k)
 84
 85    # ------------------------------------------------------------------
 86    # Core mutation
 87    # ------------------------------------------------------------------
 88
 89    def clear(self) -> None:
 90        """
 91        Remove all elements.
 92
 93        Example::
 94
 95            s = TreeMultiSet([1, 2, 3])
 96            s.clear()
 97            len(s)  # 0
 98
 99        """
100        self._t.clear()
101
102    def add(self, key: K) -> None:
103        """
104        Insert *key*, always adding a new entry even if an equal key exists.
105
106        Example::
107
108            s = TreeMultiSet()
109            s.add(1)
110            s.add(1)
111            len(s)  # 2
112
113        """
114        n: TreeNode[K, K] = TreeNode()
115        n.key = key
116        self._t.insert_multi(n)
117
118    def discard(self, key: K) -> None:
119        """
120        Remove one occurrence of *key* if present; do nothing when absent.
121
122        Example::
123
124            s = TreeMultiSet([1, 2, 2, 3])
125            s.discard(2)  # s is now {1,2,3}
126            s.discard(99)  # no-op
127
128        """
129        it = self._t.find(key)
130        if not it.equals(self._t.end()):
131            self._t.erase(it.node)
132
133    def remove(self, key: K) -> None:
134        """
135        Remove one occurrence of *key*, raising KeyError when absent.
136
137        Args:
138            key: Key to remove.
139
140        Raises:
141            KeyError: When *key* is not present.
142
143        Example::
144
145            s = TreeMultiSet([1, 2, 2, 3])
146            s.remove(2)  # s is now {1,2,3}
147            s.remove(99)  # raises KeyError
148
149        """
150        it = self._t.find(key)
151        if not it.equals(self._t.end()):
152            self._t.erase(it.node)
153        else:
154            raise KeyError(key)
155
156    def delete(self, key: K) -> None:
157        """
158        @private Remove one occurrence of *key* if present; alias for discard().
159
160        Example::
161
162            s = TreeMultiSet([1, 2, 2, 3])
163            s.delete(2)  # s is now {1,2,3}
164            s.delete(99)  # no-op
165
166        """
167        self.discard(key)
168
169    def pop(self, key: K | None = None, default: K | None = None) -> K | None:
170        """
171        Remove one occurrence of *key* and return it, or return *default* when absent.
172
173        When called with no arguments, removes and returns the smallest element.
174
175        Args:
176            key: Key to remove. When None, the smallest element is removed.
177            default: Value to return when *key* is absent.
178
179        Raises:
180            KeyError: When called with no arguments on an empty set.
181
182        Example::
183
184            s = TreeMultiSet([1, 2, 2, 3])
185            s.pop(2)  # 2, s is now {1,2,3}
186            s.pop(9)  # None
187            s.pop(9, 0)  # 0
188            s.pop()  # 1  — removes and returns smallest element
189
190        """
191        if self.size == 0:
192            raise KeyError("pop from an empty set")
193        it = self._t.find(key) if key is not None else self.begin()
194        if not it.equals(self._t.end()):
195            k = it.key
196            self._t.erase(it.node)
197            return k
198        return default
199
200    # ------------------------------------------------------------------
201    # Queries
202    # ------------------------------------------------------------------
203
204    def has(self, key: K) -> bool:
205        """
206        @private Return True when *key* exists in the set.
207
208        Example::
209
210            s = TreeMultiSet([1, 2, 3])
211            s.has(1)  # True
212            s.has(4)  # False
213
214        """
215        it = self._t.find(key)
216        return not it.equals(self._t.end())
217
218    def keys(self) -> PyIterator[K]:
219        """
220        @private Return a forward iterator over all keys in ascending order.
221
222        Example::
223
224            s = TreeMultiSet([1, 2, 2, 3])
225            list(s.keys())  # [1, 2, 2, 3]
226
227        """
228        return self._t.keys()
229
230    @property
231    def size(self) -> int:
232        """
233        @private Number of elements (counting duplicates) in the set.
234
235        Example::
236
237            s = TreeMultiSet([1, 2, 2, 3])
238            s.size  # 4
239
240        """
241        return self._t.size()
242
243    def __contains__(self, key: K) -> bool:
244        """
245        Return True when *key* is in the set.
246
247        Example::
248
249            s = TreeMultiSet([1, 2, 3])
250            1 in s  # True
251            4 not in s  # True
252
253        """
254        return self.has(key)
255
256    def __iter__(self) -> PyIterator[K]:
257        """
258        Iterate over all keys in ascending order (duplicates included).
259
260        Example::
261
262            s = TreeMultiSet([1, 2, 2, 3])
263            list(s)  # [1, 2, 2, 3]
264
265        """
266        return self._t.keys()
267
268    def __len__(self) -> int:
269        """
270        Return the total number of elements (counting duplicates).
271
272        Example::
273
274            s = TreeMultiSet([1, 2, 2])
275            len(s)  # 3
276
277        """
278        return self._t.size()
279
280    def __str__(self) -> str:
281        """
282        Return a string representation in the form {key1,key2,...}.
283
284        Example::
285
286            s = TreeMultiSet([1, 2, 2, 3])
287            str(s)  # "{1,2,2,3}"
288
289        """
290        return str(self._t)
291
292    def __repr__(self) -> str:
293        """
294        Return a string representation in the form {key1,key2,...}.
295
296        Example::
297
298            s = TreeMultiSet([1, 2, 2, 3])
299            str(s)  # "{1,2,2,3}"
300
301        """
302        return self.__str__()
303
304    # ------------------------------------------------------------------
305    # Additional iteration
306    # ------------------------------------------------------------------
307
308    def backwards(self) -> PyReverseIterator[K]:
309        """
310        Return a reverse iterator over all keys in descending order.
311
312        Example::
313
314            s = TreeMultiSet([1, 2, 3])
315            list(s.backwards())  # [3, 2, 1]
316
317        """
318        return self._t.keys().backwards()
319
320    # ------------------------------------------------------------------
321    # Custom comparator
322    # ------------------------------------------------------------------
323
324    @property
325    def compare_func(self) -> Callable[[Any, Any], int]:
326        """
327        @private The current 3-way comparison function used to order keys.
328
329        Setting this property clears all existing elements.
330
331        Example::
332
333            s = TreeMultiSet()
334            s.compare_func = lambda a, b: -1 if a < b else (1 if a > b else 0)
335
336        """
337        return self._t.compare
338
339    @compare_func.setter
340    def compare_func(self, func: Callable[[Any, Any], int]) -> None:
341        """@private Replace the comparison function and clear all elements."""
342        self.clear()
343        self._t.compare = func
344
345    # ------------------------------------------------------------------
346    # STL-style iterators
347    # ------------------------------------------------------------------
348
349    def begin(self) -> TreeIterator[K, K]:
350        """
351        Return a forward STL-like iterator to the first element.
352
353        Example::
354
355            s = TreeMultiSet([1, 2, 3])
356            it = s.begin()
357            while not it.equals(s.end()):
358                print(it.key)
359                it.next()
360            # 1, 2, 3
361
362        """
363        return self._t.begin()
364
365    def end(self) -> TreeIterator[K, K]:
366        """
367        Return a forward STL-like iterator past the last element.
368
369        Example::
370
371            s = TreeMultiSet([1, 2, 3])
372            it = s.begin()
373            while not it.equals(s.end()):
374                print(it.key)
375                it.next()
376
377        """
378        return self._t.end()
379
380    def find(self, key: K) -> TreeIterator[K, K]:
381        """
382        Return an STL-like iterator to an element with *key*, or end() if absent.
383
384        Example::
385
386            s = TreeMultiSet([1, 2, 2, 3])
387            it = s.find(2)
388            if not it.equals(s.end()):
389                print(it.key)  # 2
390
391        """
392        return self._t.find(key)
393
394    def insert_unique(self, key: K) -> InsertionResult[TreeIterator[K, K]]:
395        """
396        Insert *key* only when it is not already present.
397
398        Returns:
399            InsertionResult with was_added=True when the key was new.
400
401        Example::
402
403            s = TreeMultiSet()
404            res = s.insert_unique(1)
405            res.was_added  # True
406            res2 = s.insert_unique(1)
407            res2.was_added  # False
408            s.size  # 1
409
410        """
411        n: TreeNode[K, K] = TreeNode()
412        n.key = key
413        return self._t.insert_unique(n)
414
415    def insert_or_replace(self, key: K) -> InsertionResult[TreeIterator[K, K]]:
416        """
417        Insert *key* if absent; if present, replace in-place.
418
419        Returns:
420            InsertionResult with was_added=True on new insertion, or
421            was_replaced=True when the key already existed.
422
423        Example::
424
425            s = TreeMultiSet()
426            res = s.insert_or_replace(1)
427            res.was_added  # True
428            res2 = s.insert_or_replace(1)
429            res2.was_replaced  # True
430            s.size  # 1
431
432        """
433        n: TreeNode[K, K] = TreeNode()
434        n.key = key
435        return self._t.insert_or_replace(n)
436
437    def insert_multi(self, key: K) -> InsertionResult[TreeIterator[K, K]]:
438        """
439        Insert *key* unconditionally, even when an equal key already exists.
440
441        Returns:
442            InsertionResult always with was_added=True.
443
444        Example::
445
446            s = TreeMultiSet()
447            s.insert_multi(1)
448            s.insert_multi(1)
449            s.size  # 2
450
451        """
452        n: TreeNode[K, K] = TreeNode()
453        n.key = key
454        return self._t.insert_multi(n)
455
456    def erase(self, iterator: TreeIterator[K, K]) -> None:
457        """
458        Remove the element pointed to by *iterator*.
459
460        Example::
461
462            s = TreeMultiSet([1, 2, 3])
463            it = s.find(2)
464            it.prev()
465            s.erase(it)  # removes element with key 1
466            str(s)  # "{2,3}"
467
468        """
469        self._t.erase(iterator.node)
470
471    def lower_bound(self, key: K) -> TreeIterator[K, K]:
472        """
473        Return an STL-like iterator to the first element with key >= *key*, or end().
474
475        Example::
476
477            s = TreeMultiSet([2, 4, 4, 6])
478            lo = s.lower_bound(3)  # iterator to first 4
479            hi = s.upper_bound(4)  # iterator past last 4
480            it = lo
481            while not it.equals(hi):
482                print(it.key)  # 4, 4
483                it.next()
484
485        """
486        return self._t.lower_bound(key)
487
488    def rbegin(self) -> ReverseIterator[K, K]:
489        """
490        Return a reverse STL-like iterator to the element with the largest key.
491
492        Example::
493
494            s = TreeMultiSet([1, 2, 3])
495            it = s.rbegin()
496            while not it.equals(s.rend()):
497                print(it.key)  # 3, 2, 1
498                it.next()
499
500        """
501        return self._t.rbegin()
502
503    def rend(self) -> ReverseIterator[K, K]:
504        """
505        Return a reverse STL-like iterator before the element with the smallest key.
506
507        Example::
508
509            s = TreeMultiSet([1, 2, 3])
510            it = s.rbegin()
511            while not it.equals(s.rend()):
512                print(it.key)  # 3 2 1
513                it.next()
514
515        """
516        return self._t.rend()
517
518    def upper_bound(self, key: K) -> TreeIterator[K, K]:
519        """
520        Return an STL-like iterator to the first element with key > *key*, or end().
521
522        Example::
523
524            s = TreeMultiSet([2, 4, 4, 6])
525            lo = s.lower_bound(3)
526            hi = s.upper_bound(4)
527            it = lo
528            while not it.equals(hi):
529                print(it.key)  # 4, 4
530                it.next()
531
532        """
533        return self._t.upper_bound(key)
534
535    def first(self) -> K | None:
536        """
537        Return the smallest element, or None when the set is empty.
538
539        Example::
540
541            s = TreeMultiSet([1, 1, 2, 3])
542            s.first()  # 1
543            TreeMultiSet().first()  # None
544
545        """
546        return self._t.first()  # type: ignore[return-value]
547
548    def last(self) -> K | None:
549        """
550        Return the largest element, or None when the set is empty.
551
552        Example::
553
554            s = TreeMultiSet([1, 2, 3, 3])
555            s.last()  # 3
556            TreeMultiSet().last()  # None
557
558        """
559        return self._t.last()  # type: ignore[return-value]