[競プロ・Python] 2乗のmodを取る場合にpow関数を使うべきか

概要

Python には組み込み関数 pow があり、pow(a, b, m) で $a^b\ \mathrm{mod}\ m$ を繰り返し二乗法により高速に ( $O(\log{b})$ で) 計算してくれます。
ただ、繰り返し二乗法の高速化の恩恵を受けないほど $b$ が小さい場合、例えば $b = 2$ の場合には pow 関数を使わずに a ** 2 % m と書くこともできます。
しかし、両者は実行時間に差があり、どちらが速いかは $m$ の大きさにより異なります。

検証

n = 10 ** 7
m = 10 ** 9
# m = 10 ** 10
a = 3

for i in range(n):
    a = pow(a, 2, m)
    # a = a ** 2 % m

上記コードのコメントアウト箇所を適宜変更し、AtCoder のコードテストから PyPy3 (7.3.0) で実行
実行時間はそれぞれ $3$ 回の平均値

	$m = 10^9$	$m = 10^{10}$
pow(a, 2, m)	369 ms	397 ms
a ** 2 % m	115 ms	6163 ms

a ** 2 % m は $m = 10^9$ では速いですが、 $m = 10^{10}$ では逆転されるだけでなく大幅に遅くなります。

考察

どちらが速いかが変わっているのは、$m$ が int32 の上限 (約 $2.1 \times 10^9$) を超えると、計算途中 (a ** 2 をした瞬間) に多倍長整数 (int64 の上限を超える整数) が登場するようになることが原因と考えられます。

• $m$ が int32 の範囲内 → 関数への受け渡しなどの処理がない分、a ** 2 % m が速い
• $m$ が int32 の上限より大きい → 多倍長整数をうまく処理してくれる分、pow(a, 2, m) が速い (うまく処理できない分、a ** 2 % m が非常に遅い)

なお、(PyPy でなく) Python (3.8.2) で実行した場合は、時間に大きな差はありませんでした (4パターンいずれも数秒程度) 。
また、調べたら、多倍長整数により (特に PyPy で) 遅くなっている例がほかにも出てきました: https://inamori.hateblo.jp/entry/20130315/p1

まとめ

現在、AtCoder の問題では int32 の範囲内の数で mod をとることが多く ( $998244353$ や $10^9+7$ も範囲内です) 、$2$ 乗の場合は a ** 2 % m と書いて問題なさそうです。
一方で、特殊な $m$ が与えられる問題や、ミラーテストなどのアルゴリズムを使う際には、$m$ の大きさに注意して、必要に応じて pow(a, 2, m) を使うのが良さそうです。