Source code for titan_pylib.math.mod_array

  1from typing import Iterable, Union
  2from functools import reduce, lru_cache
  3
  4
[docs] 5class ModArray998244353: 6 7 @staticmethod 8 @lru_cache(maxsize=None) 9 def _inv(a: int) -> int: 10 res = 1 11 b = 998244351 12 while b: 13 if b & 1: 14 res = res * a % 998244353 15 a = a * a % 998244353 16 b >>= 1 17 return res 18 19 def __init__(self, n_or_a: Union[int, Iterable[int]], e: int = 0): 20 e %= 998244353 21 self.a: list[int] = ( 22 [e] * n_or_a if isinstance(n_or_a, int) else [e % 998244353 for e in n_or_a] 23 ) 24
[docs] 25 def resize(self, size: int, e: int = 0): 26 e %= 998244353 27 self.a = [e] * size
28
[docs] 29 def append(self, v: int) -> None: 30 self.a.append(v % 998244353)
31
[docs] 32 def pop(self, k: int = -1) -> int: 33 return self.a.pop(k)
34
[docs] 35 def all_sum(self) -> int: 36 return sum(self.a) % 998244353
37
[docs] 38 def all_mul(self) -> int: 39 return reduce(lambda x, y: x * y % 998244353, self.a)
40
[docs] 41 def add(self, k: int, v: int) -> None: 42 self.a[k] = (self.a[k] + v) % 998244353
43
[docs] 44 def sub(self, k: int, v: int) -> None: 45 self.a[k] = (self.a[k] - v) % 998244353
46
[docs] 47 def mul(self, k: int, v: int) -> None: 48 self.a[k] = (self.a[k] * v) % 998244353
49
[docs] 50 def div(self, k: int, v: int) -> None: 51 self.a[k] = (self.a[k] * self._inv(v)) % 998244353
52
[docs] 53 def get_mod(self) -> int: 54 return 998244353
55 56 def __iter__(self): 57 self._iter = 0 58 return self 59 60 def __next__(self): 61 if self._iter >= len(self.a): 62 raise StopIteration 63 self._iter += 1 64 return self.a[self._iter - 1] 65 66 def __getitem__(self, k: Union[int, slice]) -> Union[int, "ModArray998244353"]: 67 return ModArray998244353(self.a[k]) if isinstance(k, slice) else self.a[k] 68 69 def __setitem__(self, k: int, v: int) -> None: 70 assert isinstance(v, int) 71 self.a[k] = v % 998244353 72 73 def __str__(self): 74 return str(self.a) 75 76 def __len__(self): 77 return len(self.a)
78 79 80Array = ModArray998244353 81 82# ---------------------- # 83 84from typing import Iterable, Union 85from functools import reduce, lru_cache 86 87
[docs] 88class ModArray1000000007: 89 90 @staticmethod 91 @lru_cache(maxsize=None) 92 def _inv(a: int) -> int: 93 res = 1 94 b = 1000000005 95 while b: 96 if b & 1: 97 res = res * a % 1000000007 98 a = a * a % 1000000007 99 b >>= 1 100 return res 101 102 def __init__(self, n_or_a: Union[int, Iterable[int]], e: int = 0): 103 e %= 1000000007 104 self.a: list[int] = ( 105 [e] * n_or_a 106 if isinstance(n_or_a, int) 107 else [e % 1000000007 for e in n_or_a] 108 ) 109
[docs] 110 def resize(self, size: int, e: int = 0): 111 e %= 1000000007 112 self.a = [e] * size
113
[docs] 114 def append(self, v: int) -> None: 115 self.a.append(v % 1000000007)
116
[docs] 117 def pop(self, k: int = -1) -> int: 118 return self.a.pop(k)
119
[docs] 120 def all_sum(self) -> int: 121 return sum(self.a) % 1000000007
122
[docs] 123 def all_mul(self) -> int: 124 return reduce(lambda x, y: x * y % 1000000007, self.a)
125
[docs] 126 def add(self, k: int, v: int) -> None: 127 self.a[k] = (self.a[k] + v) % 1000000007
128
[docs] 129 def sub(self, k: int, v: int) -> None: 130 self.a[k] = (self.a[k] - v) % 1000000007
131
[docs] 132 def mul(self, k: int, v: int) -> None: 133 self.a[k] = (self.a[k] * v) % 1000000007
134
[docs] 135 def div(self, k: int, v: int) -> None: 136 self.a[k] = (self.a[k] * self._inv(v)) % 1000000007
137
[docs] 138 def get_mod(self) -> int: 139 return 1000000007
140 141 def __iter__(self): 142 self._iter = 0 143 return self 144 145 def __next__(self): 146 if self._iter >= len(self.a): 147 raise StopIteration 148 self._iter += 1 149 return self.a[self._iter - 1] 150 151 def __getitem__(self, k: Union[int, slice]) -> Union[int, "ModArray1000000007"]: 152 return ModArray1000000007(self.a[k]) if isinstance(k, slice) else self.a[k] 153 154 def __setitem__(self, k: int, v: int) -> None: 155 assert isinstance(v, int) 156 self.a[k] = v % 1000000007 157 158 def __str__(self): 159 return str(self.a) 160 161 def __len__(self): 162 return len(self.a)
163 164 165Array = ModArray1000000007 166 167# ---------------------- # 168 169from typing import Iterable, Union 170from functools import reduce, lru_cache 171 172
[docs] 173class ModArray: 174 175 @lru_cache(maxsize=None) 176 def _inv(self, a: int) -> int: 177 res = 1 178 b = self.mod - 2 179 while b: 180 if b & 1: 181 res = res * a % self.mod 182 a = a * a % self.mod 183 b >>= 1 184 return res 185 186 def __init__(self, mod: int, n_or_a: Union[int, Iterable[int]] = [], e: int = 0): 187 self.mod: int = mod 188 e %= mod 189 self.a: list[int] = ( 190 [e] * n_or_a if isinstance(n_or_a, int) else [e % mod for e in n_or_a] 191 ) 192
[docs] 193 def resize(self, size: int, e: int = 0): 194 e %= self.mod 195 self.a = [e] * size
196
[docs] 197 def append(self, v: int) -> None: 198 self.a.append(v % self.mod)
199
[docs] 200 def pop(self, k: int = -1) -> int: 201 return self.a.pop(k)
202
[docs] 203 def all_sum(self) -> int: 204 return sum(self.a) % self.mod
205
[docs] 206 def all_mul(self) -> int: 207 return reduce(lambda x, y: x * y % self.mod, self.a)
208
[docs] 209 def add(self, k: int, v: int) -> None: 210 self.a[k] = (self.a[k] + v) % self.mod
211
[docs] 212 def sub(self, k: int, v: int) -> None: 213 self.a[k] = (self.a[k] - v) % self.mod
214
[docs] 215 def mul(self, k: int, v: int) -> None: 216 self.a[k] = (self.a[k] * v) % self.mod
217
[docs] 218 def div(self, k: int, v: int) -> None: 219 self.a[k] = (self.a[k] * self._inv(v)) % self.mod
220
[docs] 221 def get_mod(self) -> int: 222 return self.mod
223 224 def __iter__(self): 225 self._iter = 0 226 return self 227 228 def __next__(self): 229 if self._iter >= len(self.a): 230 raise StopIteration 231 self._iter += 1 232 return self.a[self._iter - 1] 233 234 def __getitem__(self, k: Union[int, slice]) -> Union[int, "ModArray"]: 235 return ModArray(self.mod, self.a[k]) if isinstance(k, slice) else self.a[k] 236 237 def __setitem__(self, k: int, v: int) -> None: 238 assert isinstance(v, int) 239 self.a[k] = v % self.mod 240 241 def __str__(self): 242 return str(self.a) 243 244 def __len__(self): 245 return len(self.a)
246 247 248Array = ModArray 249 250# ---------------------- #