python technique - shakayami/ACL-for-python GitHub Wiki
テクニック集
pythonで競プロをする際に役に立つテクニック集です。これが全てであるとは限りません。
深く知りたい場合は、強い人の提出コードを見て勉強するのが良いでしょう。
pypy
提出するときにpythonではなくpypyを使うというものです。dpなどはこれで速くなります。 注意点としては、numpy等を使っている場合はREになります。
dfs等で再帰を多く使っている場合は、pypyよりpythonの方が速くなる場合があります。
pypyだとMLEになるコードがpythonだとACになる場合があります。
numpy
pypyは使えなくてpythonでの話になりますが、numpyを使うと高速化できる場合があります。numpyには機能がたくさんあるので書ききれないのですが例を一つ。 畳み込みは普通に書くよりも定数倍高速化が出来ます。
import numpy as np
K=1<<20
N,M=map(int,input().split())
a=list(map(int,input().split()))
b=list(map(int,input().split()))
A=a+[0]*(K-N)
B=b+[0]*(K-M)
C=np.fft.ifft(np.fft.fft(B)*np.fft.fft(A))
ans=[int(i+0.5) for i in np.real(C)][:N+M-1]
print(*ans)
注意点としては、ACL practice contest のF-Convolutionはこれを使えません。ACLでいうとF問題はConvolutionを使っているのに対して、これはConvolution-llであるからです。
deque
import Queue
で定義されるキューを使うよりも、
from collections import deque
で定義されるキューを使うと高速になります。 しかもこのキューは両端キューであるため、（つまり、キューとスタックの両方の性質を併せ持っている）機能面でも優れています。
入力
import sys
input=sys.stdin.readline
おまじないみたいなものです。
ただし、文字列を入力として受け取る場合にこれを使うと、改行コードである"\n"までも入力に含んでしまうため、WAの原因となる可能性があります。
再帰制限
import sys
sys.setrecursionlimit(10**6)
dfsなどで再帰するとき、再帰的なプログラミングだと、再帰制限回数（デフォルトでは1000回）を超えてしまいREとなってしまいます。それを防ぐために、再帰制限回数を設定で変更しているのです。 ACL-for-pythonでは、two-satとscc（まれにmaxflow）を使う際にこの「おまじない」が無いとREになる可能性があります。 setrecursionlimitに設定できる最大の数は2147483647(=2^31-1)です。
mod逆元
C++のライブラリではmodの逆元を計算しているときに少々長めのコードを使っています。これをpythonで翻訳して書き写すのもいいですが、結局の所以下のように書くのが最も楽で高速になります。
mod=998244353
a=int(input())
ans=pow(a,mod-2,mod)
print(ans)
ただしこれはmodが素数であるときにだけ使えます。
また、pythonの新しいバージョンの場合だと以下のような表記もできます。
mod=998244353
a=int(input())
ans=pow(a,-1,mod)
print(ans)
これはaとmodが互いに素なら使えます。modが素数である必要がありません。 自分が確認した場合だと（2020年10月16日現在）、AtCoderの場合なら
- 言語アップデート後のpythonなら通過する
- pypyの場合REになる
といった感じでした。
階乗の前計算
階乗とその逆元を前計算するときには、逆元の計算は一回で済みます。以下のように書きます。
mod=998244353
MAX_N=10**5
Fact=[0]*(MAX_N+1)
Finv=[0]*(MAX_N+1)
Fact[0]=1
for i in range(MAX_N):
Fact[i+1]=(Fact[i]*(i+1))%mod
Finv[MAX_N]=pow(Fact[MAX_N],mod-2,mod)
for i in range(MAX_N-1,-1,-1):
Finv[i]=((i+1)*Finv[i+1])%mod
#Fact[i]:i! のmod
#Finv[i]:1/i! のmod
#二項係数
def C(n,k):
return ((Fact[n]*Finv[n-k])%mod*Finv[k])%mod
このようにすると計算量はO(MAX_N+log(mod))となります。これは、1/i!=(i+1)*1/(i+1)!を利用しています。 一方、例えば以下のように書くと遅くなります。
for i in range(MAX_N+1):
Finv[i]=pow(Fact[i],mod-2,mod)
このとき計算量はO(MAX_N log(mod))となります。mod=998244353ならlog(mod)は30くらいになります。30倍計算量が増えるというのは意外と厄介で、MAX_N=10^6程度のときにTLEの原因になることがあります。上記のような書き方をマスターすると良いことがあります。
また、二項係数を以下のように書くと遅くなる場合があります。
def C(n,k):
return (Fact[n]*Finv[n-k]*Finv[k])%mod
何故かと言うと、計算途中で大きな数になることで、掛け算に時間がかかるからです。定数倍のためこれで引っかかる場合があります。
集合とビット演算
集合を扱う場合、全体集合の範囲が小さいならばbitとして集合を扱ったほうが良い場合があります。
例えばABC147-Eは、dp[i][j]に「経路の総和がkになることがあるかどうか」という集合を持っています。kの範囲としては161 * 161 * 2までの範囲を持てば良いのですが、これはsetとしてデータを持つより「kビット目が1⇔経路の総和の差がkになることがあるか」となるような整数をdpテーブルに持つほうが高速になります。コードは以下のようになります。
(https://atcoder.jp/contests/abc147/tasks/abc147_e)
H,W=map(int,input().split())
A=[[int(i) for i in input().split()] for j in range(H)]
B=[[int(i) for i in input().split()] for j in range(H)]
dp=[[0 for j in range(W)] for i in range(H)]
M=161*161
dp[0][0]|=1<<(A[0][0]-B[0][0]+M)
dp[0][0]|=1<<(-A[0][0]+B[0][0]+M)
for i in range(H):
for j in range(W):
dij=abs(A[i][j]-B[i][j])
if i>0:
dp[i][j]|=dp[i-1][j]<<dij
dp[i][j]|=dp[i-1][j]>>dij
if j>0:
dp[i][j]|=dp[i][j-1]>>dij
dp[i][j]|=dp[i][j-1]<<dij
ans=[]
for k in range(2*M+1):
if (1<<k)&dp[-1][-1]!=0:
ans.append(abs(k-M))
print(min(ans))
順序付き集合
pythonにはC++でいうstd::setに対応するものが無いみたいです。
pythonにもsetはありますが、順序がついてないので二分探索ができません。
よって、std::setを想定解とした問題ではpython勢が不利になる場合があります。
ということでライブラリの紹介です。これはACL-for-Pythonの作者がつくったものではありません。
https://github.com/grantjenks/python-sortedcontainers
ライセンスについては、これはApache Licence 2.0で公開されているため、競技プログラミングに使う分には問題がないと思います。
他には、tatyam氏が作成したSortedSetもあります。