0. はじめに

競技プログラミングでは「答えを〇〇で割ったあまりを求めてください。」という問題がしばしば出題され、このようなタイプの問題で多くの参加者が使う便利なものとして ModInt というものがあります。

ModInt は int 型のように整数型として振る舞いつつ、自動であまりを取ってくれます。
例えば、$50!$ を $97$ で割ったあまりを求める場合、特に固定長整数型を使う場合はオーバーフローに注意して

mod = 97
result = 1
for i in range(1, 51):
    result *= i
    result %= mod

のように書きますが、ModInt を用いると

mod = 97
ModInt.set_mod(mod)
result = ModInt(1)
for i in range(1, 51):
    result *= i

のようにあまりをとる操作を書かなくてもよくなります。
また、割り算も、a * inv_mod(b) のようにせず、 a / b と書くことができます。

このように非常に便利なので、検索すると Python での実装例も（以下の私の記事も含めて）たくさん出てきます。

【Modint編】AtCoder Library 解読〜Pythonでの実装まで〜

しかし、Python で実装した場合オーバーヘッドが非常に大きく、実用的な(コンテストで使えるほどの)速度が出ないことに加えて、組み込み関数pow()の存在も後押しし、「手動で mod をとればいいか」という気持ちになる人が多いと思います。

そこで、使ってみたいと思えるようなものを作るべく、Python/C API を使って ModInt を実装してみました。

1. Python/C API

Python/C API を使うことで、 C/C++ で書かれたモジュールを Python から利用できるようになります。

これによって「組み込み型の int + 手動 mod」という方法と同程度の速度を実現する ModInt をつくろうという考えです。

2. 実装のお手本

標準的なPython (CPython) はC言語で実装されており、そのソースコードは https://github.com/python/cpython で見ることができます。
その中には、組み込み型の int に相当する PyLongObject の実装もあります。
また、ModInt に関しては AtCoder が出している AtCoder Library の実装を見ることができます。

これらをいい具合に混ぜ合わせることで、ModInt for Python をつくってみました。

3. コード

https://github.com/Akari-Tk/ac-library-py/blob/main/atcoder/modint.hpp

コードはこちらで見ることができます。

4. ベンチマーク

以下のコードで「組み込み型の int + 手動 mod」と今回実装した ModInt を比較しました。

import time
from atcoder import ModInt as mint


def benchmark(func):
    def _wrapper(*args, **keywargs):
        t0 = time.perf_counter()
        result = func(*args, **keywargs)
        txt = f'[{func.__name__}] time: {time.perf_counter()-t0:6.3f} sec'
        if result is not None: txt += f'{result = }'
        print(txt)
        return result
    return _wrapper




@benchmark
def add_mint(n):
    a = mint(0)
    for i in range(n):
        a += i

@benchmark
def sub_mint(n):
    a = mint(0)
    for i in range(n):
        a -= i

@benchmark
def mul_mint(n):
    a = mint(1)
    for i in range(1, n):
        a *= i

@benchmark
def div_mint(n):
    a = mint(1)
    for i in range(1, n):
        a //= i

@benchmark
def pow_mint(n):
    a = mint(2)
    for i in range(n):
        a ** i


@benchmark
def add__int(n, mod):
    a = 0
    for i in range(n):
        a += i
        a %= mod

@benchmark
def sub__int(n, mod):
    a = 0
    for i in range(n):
        a -= i
        a %= mod

@benchmark
def mul__int(n, mod):
    a = 1
    for i in range(1, n):
        a *= i
        a %= mod

@benchmark
def div__int(n, mod):
    a = 1
    for i in range(1, n):
        a *= pow(i, -1, mod)
        a %= mod

@benchmark
def pow__int(n, mod):
    a = 2
    for i in range(n):
        pow(a, i, mod)




def main():
    mod = 998244353
    mint.set_mod(mod)
    n = 10_000_000
    print('\n---Benchmark---')
    print(f'mod = {mod:,}')
    print(f'  n = {n:,}\n')
    add_mint(n)
    add__int(n, mod)
    sub_mint(n)
    sub__int(n, mod)
    mul_mint(n)
    mul__int(n, mod)
    div_mint(n)
    div__int(n, mod)
    pow_mint(n)
    pow__int(n, mod)


if __name__ == '__main__':
    main()

結果は以下のようになりました。どの演算も手動 mod より高速になり、掛け算、割り算、累乗では 2〜16倍高速 な ModInt ができました。

---Benchmark---
mod = 998,244,353
  n = 10,000,000

[add_mint] time:  0.685 sec
[add__int] time:  0.877 sec
[sub_mint] time:  0.673 sec
[sub__int] time:  0.893 sec
[mul_mint] time:  0.673 sec
[mul__int] time:  1.342 sec
[div_mint] time:  2.095 sec
[div__int] time: 19.916 sec
[pow_mint] time:  1.676 sec
[pow__int] time: 27.801 sec

5. AtCoder への提出

AtCoder では Python (3.8.2) を選択するとコンパイルフェーズがあるのでこれを利用することで提出できます。

実際にトヨタ自動車プログラミングコンテスト2023#1(AtCoder Beginner Contest 298) の D問題『Writing a Numeral』に提出してみました。

提出

実行時間制限が $3 \:\rm{s}$ の問題ですが $543 \:\rm{ms}$ で AC することができました。
これは提出時点における Python (3.8.2) での AC の中で 3 番目の速度です。

6. 注意

このような C拡張を埋め込む方法はコンテストサイトによってはルール違反となりアカウントBANなどの措置が取られることがあります。特に、Codeforces では禁じられているので注意してください。(ルール)
どのコンテストサイトであっても提出する前にコンテストルールを確認してください。
また、紹介したものはバグを含む可能性がありますので使用する場合は自己責任でお願いします。(バグがあったら報告いただけると嬉しいです。)

7. おわりに

Python で実用的な ModInt を使うために Python/C API に手を出してみました。
Python の実装について色々わかったので面白かったです。

バグ、改善点などありましたらお教えください。

Python/C API を利用して高速な ModInt をつくってみた