diff --git a/library/modulos.py b/library/modulos.py new file mode 100644 index 0000000..563e246 --- /dev/null +++ b/library/modulos.py @@ -0,0 +1,246 @@ +# originally written on 2019/05/21 ? +# added division on 2020/06/19 + +def inverse(a, modulo): + def is_even(a): + return a & 1 == 0 + + def rotate_right(a): + return a >> 1 if is_even(a) else (a + modulo) >> 1 + + if a == 0: + raise ZeroDivisionError() + b = 1 + c = 0 + d = modulo + + while a != d: + if a < d: + d -= a + c = (c - b) % modulo + while is_even(d): + d >>= 1 + c = rotate_right(c) + else: + a -= d + b = (b - c) % modulo + while is_even(a): + a >>= 1 + b = rotate_right(b) + + return b + +class ModInt: + def __init__(self, value, modulo): + assert isinstance(modulo, int), (type(modulo), modulo) + assert modulo >= 0 + if isinstance(value, ModInt): + if value.modulo == modulo: + self.value = value.value + else: + self.value = value.value % modulo + else: + assert isinstance(value, int), (type(value), value) + self.value = value % modulo + self.modulo = modulo + + def new(self, value): + return ModInt(value, self.modulo) + + @staticmethod + def as_ints(a, b, op="?"): + if isinstance(a, ModInt): + if isinstance(b, ModInt): + assert a.modulo == b.modulo, \ + f"mismatched modulos: {repr(a)} {op} {repr(b)}" + return a.value, b.value + else: + assert isinstance(b, int), \ + f"incompatible types: {repr(a)} {op} {b} ({type(b)})" + return a.value, b + elif isinstance(b, ModInt): + assert isinstance(a, int), \ + f"incompatible types: {repr(a)} ({type(a)}) {op} {repr(b)}" + return a, b.value + else: + assert isinstance(a, int) and isinstance(b, int), ( + "incompatible types: " + + f"{repr(a)} ({type(a)}) {op} {repr(b)} ({type(b)})") + return a, b + + def __add__(self, other): + a, b = self.as_ints(self, other, "+") + return self.new(a + b) + + def __radd__(self, other): + a, b = self.as_ints(other, self, "+") + return self.new(a + b) + + def __sub__(self, other): + a, b = self.as_ints(self, other, "-") + return self.new(a - b) + + def __rsub__(self, other): + a, b = self.as_ints(other, self, "-") + return self.new(a - b) + + def __mul__(self, other): + a, b = self.as_ints(self, other, "*") + return self.new(a * b) + + def __rmul__(self, other): + a, b = self.as_ints(other, self, "*") + return self.new(a * b) + + def __truediv__(self, other): + a, b = self.as_ints(self, other, "/") + return self.new(a * inverse(b, self.modulo)) + + def __rtruediv__(self, other): + a, b = self.as_ints(other, self, "/") + return self.new(a * inverse(b, self.modulo)) + + def __floordiv__(self, other): + a, b = self.as_ints(self, other, "//") + return self.new(a // b) + + def __rfloordiv__(self, other): + a, b = self.as_ints(other, self, "//") + return self.new(a // b) + + def __mod__(self, other): + a, b = self.as_ints(self, other, "%") + return self.new(a % b) + + def __rmod__(self, other): + a, b = self.as_ints(other, self, "%") + return self.new(a % b) + + def __divmod__(self, other): + a, b = self.as_ints(self, other, "divmod") + return self.new(a // b), self.new(a % b) + + def __pow__(self, other, modulo=None): + a, b = self.as_ints(self, other, "**") + if modulo is None: + return self.new(pow(a, b, self.modulo)) + else: + return self.new(pow(a, b, modulo)) + + def __rpow__(self, other): + a, b = self.as_ints(other, self, "**") + return self.new(pow(a, b, self.modulo)) + + def __eq__(self, other): + a, b = self.as_ints(self, other, "==") + return a == b + + def __ne__(self, other): + a, b = self.as_ints(self, other, "!=") + return a != b + + def __gt__(self, other): + a, b = self.as_ints(self, other, ">") + return a > b + + def __lt__(self, other): + a, b = self.as_ints(self, other, "<") + return a < b + + def __ge__(self, other): + a, b = self.as_ints(self, other, ">=") + return a >= b + + def __le__(self, other): + a, b = self.as_ints(self, other, "<=") + return a <= b + + def __iadd__(self, other): + a, b = self.as_ints(self, other, "+=") + self.value += b + self.value %= self.modulo + return self + + def __isub__(self, other): + a, b = self.as_ints(self, other, "-=") + self.value -= b + self.value %= self.modulo + return self + + def __imul__(self, other): + a, b = self.as_ints(self, other, "*=") + self.value *= b + self.value %= self.modulo + return self + + def __itruediv__(self, other): + a, b = self.as_ints(self, other, "/=") + self.value *= inverse(b) + self.value %= self.modulo + return self + + def __ifloordiv__(self, other): + a, b = self.as_ints(self, other, "//=") + self.value //= b + self.value %= self.modulo + return self + + def __imod__(self, other): + a, b = self.as_ints(self, other, "%=") + self.value %= b + self.value %= self.modulo + return self + + def __ipow__(self, other): + a, b = self.as_ints(self, other, "**=") + self.value = pow(a, b, self.modulo) + return self + + def __pos__(self): + return self + + def __neg__(self): + return 0 - self + + def __abs__(self): + return self + + def __invert__(self): + raise NotImplementedError("operator unimplemented for ModInt: ~") + + def __int__(self): + return self.value + + def __float__(self): + return float(self.value) + + def __complex__(self): + return complex(self.value) + + def __index__(self): + return self.value + + def __round__(self, digits=None): + return self + + def __trunc__(self): + return self + + def __floor__(self): + return self + + def __ceil__(self): + return self + + def __repr__(self): + return f"{self.__class__.__name__}({self.value}, {self.modulo})" + + def __str__(self): + return str(self.value) + + def __hash__(self): + # it's an error to use ModInts alongside others of differing modulo, + # so this should just act like an int here. + # (hmm, maybe it's better to return (self.value, self.modulo) anyway. + # or just leave it unhashable. then the user has to be conscious.) + return self.value