Source code for titan_pylib.data_structures.set.min_max_multiset

  1from titan_pylib.my_class.supports_less_than import SupportsLessThan
  2from titan_pylib.data_structures.heap.double_ended_heap import DoubleEndedHeap
  3from typing import Generic, Iterable, TypeVar
  4
  5T = TypeVar("T", bound=SupportsLessThan)
  6
  7

[docs]
  8class MinMaxMultiset(Generic[T]):
  9
 10    def __init__(self, a: Iterable[T] = []):
 11        a = list(a)
 12        data = {}
 13        for x in a:
 14            if x in data:
 15                data[x] += 1
 16            else:
 17                data[x] = 1
 18        self.data = data
 19        self.heap = DoubleEndedHeap(a)
 20        self.len = len(a)
 21

[docs]
 22    def add(self, key: T, val: int = 1) -> None:
 23        if val == 0:
 24            return
 25        self.heap.add(key)
 26        if key in self.data:
 27            self.data[key] += val
 28        else:
 29            self.data[key] = val
 30        self.len += val

 31

[docs]
 32    def discard(self, key: T, val: int = 1) -> bool:
 33        if key not in self.data:
 34            return False
 35        cnt = self.data[key]
 36        if val < cnt:
 37            self.len -= val
 38            self.data[key] -= val
 39        else:
 40            self.len -= cnt
 41            del self.data[key]
 42        return True

 43

[docs]
 44    def pop_min(self) -> T:
 45        while True:
 46            v = self.heap.get_min()
 47            if v in self.data:
 48                if self.data[v] == 1:
 49                    self.heap.pop_min()
 50                    del self.data[v]
 51                else:
 52                    self.data[v] -= 1
 53                self.len -= 1
 54                return v
 55            self.heap.pop_min()

 56

[docs]
 57    def pop_max(self) -> T:
 58        while True:
 59            v = self.heap.get_max()
 60            if v in self.data:
 61                self.len -= 1
 62                if self.data[v] == 1:
 63                    self.heap.pop_max()
 64                    del self.data[v]
 65                else:
 66                    self.data[v] -= 1
 67                return v
 68            self.heap.pop_max()

 69

[docs]
 70    def get_min(self) -> T:
 71        while True:
 72            v = self.heap.get_min()
 73            if v in self.data:
 74                return v
 75            else:
 76                self.heap.pop_min()

 77

[docs]
 78    def get_max(self) -> T:
 79        while True:
 80            v = self.heap.get_max()
 81            if v in self.data:
 82                return v
 83            else:
 84                self.heap.pop_max()

 85

[docs]
 86    def count(self, key: T) -> int:
 87        return self.data[key]

 88

[docs]
 89    def tolist(self) -> list[T]:
 90        return sorted(k for k, v in self.data.items() for _ in range(v))

 91

[docs]
 92    def len_elm(self) -> int:
 93        return len(self.data)

 94
 95    def __contains__(self, key: T):
 96        return key in self.data
 97
 98    def __len__(self):
 99        return self.len
100
101    def __str__(self):
102        return "{" + ", ".join(map(str, self.tolist())) + "}"
103
104    def __repr__(self):
105        return "MinMaxMultiset([" + ", ".join(map(str, self.tolist())) + "])"