backyard/library/modulos.py

# 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