Pythonでmodintを実装してみた

# はじめに

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

ちなみに、 実装したきっかけは、[エクサウィザーズ2019 E - Black or White](https://atcoder.jp/contests/exawizards2019/tasks/exawizards2019_e)で逆元の計算や `%` 演算子の嵐に悩まされたことです。

# modintのコード

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

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

```python
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$ 逆元）。

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

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

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

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

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

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

## 例3: sum関数

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

```python
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 で割ったあまり」の求め方を総特集！ 〜 逆元から離散対数まで 〜](https://qiita.com/drken/items/3b4fdf0a78e7a138cd9a) が参考になります。

ちなみに、 `__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](https://atcoder.jp/contests/exawizards2019/tasks/exawizards2019_e)  に対して、 `ModInt` を用いて以下のソースコードで解きました（PyPy3 でないとTLEになるので注意）。

https://atcoder.jp/contests/exawizards2019/submissions/4790083

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