Pythonでmodintを実装してみた

はじめに

競技プログラミング界隈で流行っているmodint（intと同じ感覚で演算可能であり、かつ、自動的に mod を取ってくれるデータ型）をPythonで実装しました。コードと実装メモを載せています。
C++での実装は、noshi91さんの記事 modint 構造体を使ってみませんか？ (C++) が参考になります。

ちなみに、 実装したきっかけは、エクサウィザーズ2019 E - Black or Whiteで逆元の計算や % 演算子の嵐に悩まされたことです。

modintのコード

ModInt クラスとして定義しました。
modulo（余り）は定数 MOD として最初に定義しています。
注意: MOD は素数である必要があります。

修正(2019/4/1 18:23): @shiracamus さんのコメントを参考に __repr__ の定義を変更しました。

class ModInt:
    def __init__(self, x):
        self.x = x % MOD

    def __str__(self):
        return str(self.x)

    __repr__ = __str__

    def __add__(self, other):
        return (
            ModInt(self.x + other.x) if isinstance(other, ModInt) else
            ModInt(self.x + other)
        )

    def __sub__(self, other):
        return (
            ModInt(self.x - other.x) if isinstance(other, ModInt) else
            ModInt(self.x - other)
        )

    def __mul__(self, other):
        return (
            ModInt(self.x * other.x) if isinstance(other, ModInt) else
            ModInt(self.x * other)
        )

    def __truediv__(self, other):
        return (
            ModInt(
                self.x * pow(other.x, MOD - 2, MOD)
            ) if isinstance(other, ModInt) else
            ModInt(self.x * pow(other, MOD - 2, MOD))
        )

    def __pow__(self, other):
        return (
            ModInt(pow(self.x, other.x, MOD)) if isinstance(other, ModInt) else
            ModInt(pow(self.x, other, MOD))
        )

    __radd__ = __add__

    def __rsub__(self, other):
        return (
            ModInt(other.x - self.x) if isinstance(other, ModInt) else
            ModInt(other - self.x)
        )

    __rmul__ = __mul__

    def __rtruediv__(self, other):
        return (
            ModInt(
                other.x * pow(self.x, MOD - 2, MOD)
            ) if isinstance(other, ModInt) else
            ModInt(other * pow(self.x, MOD - 2, MOD))
        )

    def __rpow__(self, other):
        return (
            ModInt(pow(other.x, self.x, MOD)) if isinstance(other, ModInt) else
            ModInt(pow(other, self.x, MOD))
        )

使用例

ModInt のmodulo（余り）は素数 $M$ とします。

例1: 除算

標準入力から二つの整数 $A$, $B$ を受け取り、 $AB^{-1} \mod M$ を標準出力する例です（$B^{-1}$ は $M$ における $B$ 逆元）。

A, B = map(ModInt, (map(int, input().split())))
print(A / B)

この例のように、 (A * inv(B)) % MOD と書かなくて済みます。

例2: （通常の）intとの演算

標準入力から単一の整数 $A$ を受け取り、 $(1 + A)^3 \mod M$ を標準出力する例です。

A = ModInt(int(input()))
print((1 + A)**3)

この例のように、 int と ModInt の演算が可能です。また、 int の値は左右どちらでも構いません。

例3: sum関数

標準入力から単一の整数 $N$ を受け取り、 $\sum_{i=1}^{N} i^2 \mod M$ を標準出力する例です。

N = int(input())
print(sum(ModInt(i)**2 for i in range(1, N + 1)))

この例のように、 sum 関数を適用可能です。

実装メモ

ModInt の各メソッドの実装について説明します。

init

Pythonでは a % MOD の値は、たとえ a が負であっても、 0 以上 MOD 未満の値になります。
そのため、競技プログラミングでは都合がよく、素直にコンストラクタに渡された値に対して MOD で割った余りをとっています。

str, repr

__str__ は str 関数を適用したときに int と同一の文字列が得られるように定義しています。
__repr__ は print 関数を適用したときに int と同一の文字列が標準出力されるように定義しています。

add, sub, mul

__xxx__ は演算子のオーバーライドです。例えば、 __add__ は + 演算子のオーバーライドになります。

isinstance(other, ModInt) は other が ModInt （のインスタンス）であるかを判定します。
これにより、以下を実現しています。

• 右側の被演算子 other が ModInt ならば other から整数を取り出して演算します。
• そうでなければ other 自体と演算します。

truediv

$\mod M$ における逆元を求めて、逆元との積で実現しています。
逆元はフェルマーの小定理に基づいて求めます。
組み込み関数 pow は第3引数にmodulo（余り）を指定すると、余りをとる前の値が巨大な場合でも、余りをとった後の値を高速に計算してくれます。
フェルマーの小定理に基づいた求め方は 「1000000007 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜 が参考になります。

ちなみに、 __div__ ではないことに注意してください。
Pythonには除算に関する演算子として / と // がありますが、前者が __truediv__, 後者が __floordiv__ に対応します。

pow

単に (self.x ** other.x) % MOD と書くと (self.x ** other.x) が巨大な場合に大きな計算時間がかかることに注意してください。
この問題に対しては、組み込みの pow 関数の第3引数に MOD を指定することで回避しています。

radd, rmul

接頭辞がrであるメソッドを定義することで 1 + ModInt(2) のように int, ModInt の順序での演算が可能となります。
もし、 これらのメソッドを定義しない場合、 ModInt(2) + 1 は ModInt(3) になるが、 1 + ModInt(2) はエラーになることに注意してください。
交換則が成り立たないのは事故のもとになりやすいので、定義しておいたほうが無難だと思います

rsub, rtruediv, rpow

これも __radd__, __rmul__ と同様の理由で定義したほうがいいのですが、引数の順序に注意してください。
例えば、 2 - ModInt(1) を評価すると ModInt.__rsub__(1, 2) が呼び出されます。
当然といえば当然ですが、 self には ModInt の値が入ります。
したがって、 __sub__, __truediv__, __pow__ と比べて self と other を逆にして計算する必要があります。

おわりに

とくに $\mod M$ における逆元の計算が必要な問題では、計算式に本質的でない演算が多く含まれるため、コードが汚くなりがちです。
実装の事故率を減らすという意味でもmodintを定義し使用すると幸せになれると思います。

ちなみに、私がmodintを実装するきっかけとなった問題 エクサウィザーズ2019 E - Black or White に対して、 ModInt を用いて以下のソースコードで解きました（PyPy3 でないとTLEになるので注意）。

結構直感的なコードにできたかなと思っています。