LoginSignup
50
34

More than 3 years have passed since last update.

Pythonでmodintを実装してみた

Last updated at Posted at 2019-03-31

はじめに

競技プログラミング界隈で流行っている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)

この例のように、 intModInt の演算が可能です。また、 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)otherModInt (のインスタンス)であるかを判定します。
これにより、以下を実現しています。

  • 右側の被演算子 otherModInt ならば 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) + 1ModInt(3) になるが、 1 + ModInt(2) はエラーになることに注意してください。
交換則が成り立たないのは事故のもとになりやすいので、定義しておいたほうが無難だと思います

rsub, rtruediv, rpow

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

おわりに

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

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

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

50
34
4

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up
50
34