backyard/library/modulos.py

246 lines
6.6 KiB
Python

# 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