library: add modulos.py

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