普段競技プログラミングでPythonを使っているのですが、想定解（？）でもTLEになる場合があります。その時は定数倍で高速化するらしいのですが、ちょっと特殊（？）なパターンに遭遇したので備忘録。

結論

Pythonのintはデフォルトで28byteだよ
これを超えるとメモリを再確保するよ
毎回やってるとおっそいよ

詳細

やらかした問題はTypical DP ContestのE-数です。典型的な桁DPで、

dp[i][smaller][mod] := ( i 番目まで選んだ時の答え(Dの倍数) )
ただし、smallerNより小さいかどうかのフラグ、modはDで割ったときの余りを示す。

とすれば解けます。

自分が提出したソースコードはこんな感じ。

MOD = 1000000007  
D = int(input())
N = [int(x) for x in list(input())]
dp = [[[ 0 for _ in range(D) ] for _ in range(2) ] for _ in range(len(N)+1)]

dp[0][0][0] = 1
for i in range(len(N)):
    max_lim = int(N[i])
    for l in range(2):
        for m in range(D):
            lim = 9 if l else N[i]
            lim = 9 if l else max_lim
            for d in range(lim+1):
                if l:
                    dp[i+1][l][(m+d) % D] += dp[i][l][m]
                elif d < lim:
                    dp[i+1][1][(m+d) % D] += dp[i][l][m]
                else:
                    dp[i+1][l][(m+d) % D] += dp[i][l][m]

print((dp[len(N)][0][0] + dp[len(N)][1][0] -1) % MOD)

オーダー的にはO(nD)(nはNの桁数)になるので、問題ないはずです。が、制約上限の値を入力するとPyPyで提出してもTLEになってしまいます。しょうがなくC++で再実装を行っているときに、オーバーフローを考慮した実装を行ったことでようやく気づきました…。巨大な数値をそのまま扱っているせいだ、と。

Pythonのintはデフォルトで28Byteですが、メモリが許す限り拡張し続けます。演算するたびに新しくメモリを確保…を繰り返していたらそりゃ遅いです。なので、dpを更新するたびにmodを取ることにしました。修正後のコードはこんな感じ（一部抜粋）

(~~略~~)
for d in range(lim+1):
if l:
    dp[i+1][l][(m+d) % D] += dp[i][l][m]
        dp[i+1][l][(m+d) % D] %= MOD
elif d < lim:
    dp[i+1][1][(m+d) % D] += dp[i][l][m]
        dp[i+1][l][(m+d) % D] %= MOD
else:
    dp[i+1][l][(m+d) % D] += dp[i][l][m]
        dp[i+1][l][(m+d) % D] %= MOD
(~~略~~)

TLEだったのが最大300ms強になりAC。ちなみにローカルで実行したら半分くらいの時間になりました。

研究でも競プロでも個人的な開発でもPythonを多用しているので、制限なしintに甘えてましたね……今後気を付けたいと思います。

しまりすのしっぽ。

競プロ(Python)で巨大な数を扱うときは注意しよう（自戒）

結論

詳細