じゅっぴーダイアリー

雑談・競プロ参加記・ライブラリ残し Twitter: @juppyjappy

Python 競技プログラミング高速化tips (PythonでAtcoderをやる際に個人的に気を付けてること)

こんにちは。最近やよい軒の彩定食にハマってるじゅっぴーです。
 
自分の確認と最近Python競技プログラミング始めたよーという人向けを兼ねたPython高速化記事です。
競技プログラミングAtcoderを想定しています。

はじめに


 
最近PythonAtcoderをはじめている人がどんどん増えています。
一方でPythonの高速化テクニック:C++で書き直す。』というネタがあるほど、Pythonは劇遅です。
競技プログラミングに向いていません。
最上位層(橙・赤)にPythonで行くのは多分無理です。
(2019/11/4 追記: 今日のAGCでそんなことないことがmaspyさんによって証明されました。
どうやら行けるようです。すみませんでした)
 
「まぁ、でもとりあえずPythonでいけるところまでいこー」って人がいるとして、これはそういう人向けのPython高速化Tipsの記事です。
 
今のところ自分は青ですが、Atcoder公式のRatedコンテストで「Pythonだから解けなかった、、」みたいなことは特にないです。
 
知っている内容や何を言ってるんだこやつ?って内容は読み飛ばしてどんどん次に行って全然OKです。
 
またスクリプト言語に対する言及をしている方は、一度この記事に目を通して頂けると幸いです。
スクリプト言語などで競プロをすることについて - ARMERIA
 
 

Pypyを使う!

何よりも最初に着手すべきなのは「Pypyで提出する!」というものです。
C++で書き直さなくてもPypy3として提出すると通るケースがかなりあります。
 
「この問題Pythonでの正解者ほぼいないけど、Python使ってる人壊滅したんか、、、」ってときは大体みんなPypyで提出しています。

f:id:jupy:20190611203020p:plain
下の方にあるこれ!!!
 
 
pypy.org
 
基本的にPython同じコードで提出できて、かつPythonより高速です。
 
例えば以下のようなコードをAtcoderのコードテストで実行すると、

n = 10**7
a = 0
for i in range(n):
  a = a+1

 Python3(3.4.3) : 952 ms
 Pypy3(2.4.0) : 227 ms

くらい違います。
 
「じゃあ全部Pypyで出せばよくない?」ってなると思うのですが、実際割とそうです。明らかにTLEしない問題はホイホイPythonで出したりもします。
でもPypyが苦手なこともあるので、それは後ろの方で書きます。
 
 
 

みんな一度は通る道

今回の記事の趣旨とは逸れますが、誰もが一度は通る道なので書いておきます。
 
math.gcdはだめ!( ;∀;)
Python3.4.3では最大公約数を求めるgcdは”math"ではなくて、"fractions"の中にあります。
そのため、

from fractions import gcd

とする必要があります。
 
del,insert
delやinsertはO(n)です。なので、基本使いません。
この辺りがあいまいな人は下の記事に目を通しておくといいです。
Pythonistaなら知らないと恥ずかしい計算量のはなし - Qiita
 
copy
みんな大好きなコレです。

a = [0,1,2,3]
b = a
b[1] = 4
print(a)
#>>[0, 4, 2, 3]

どうしても配列のコピーを使いたいときはdeepcopyすればいいですが、大抵の場合工夫するとこういう処理をしなくてもなんとかなります。
 
再帰限界
Pythonでは再帰限界が1000回まで決められているので、再帰を使う際はこの限界を上げてあげる必要があります。
具体的には以下のようなものを冒頭に書きます。

import sys
sys.setrecursionlimit(4100000)

そもそもPython再帰がやや遅いのであまり使わない方がいいという話もあります。
 
 
 

Pypy一択なもの

問題によってはPypy一択なものも存在します。例えば、オーダーが10^7程度となるとほとんどの場合Pypyで出すことになります。
またnumpyをうまく使わない限り、Pythonの配列計算(配列アクセス)は遅いです。
そのため、以下のような問題は基本迷わずPypyです。
 
・10^7のオーダー
・H×W盤面など とにかくマス(配列)でゴリゴリする系
・DP
・木問題
・グラフ問題

 
 

Pypyじゃだめなもの

*1
再帰
ありえないくらい遅い(Pythonもやや遅い)。再帰を深くすると完全にアウトになることがあります。
Pythonも同じなのですが、基本的には関数を呼び出するよりそのまま書く方が高速です。

defaultdict
defaultdictもやや遅いです。実装コストが大きく変わるわけではないので、dictを使うといいと思います。
またPypyでは、elt in d.keys()よりelt in dの方がはるかに速いです。これはAtcoderのver.のPypyではdictviewに__contains__がないためです。

文字関連
この世の終わりみたいな速度。実質 sys.exit(0)

n = 10**5
res = 'a'
for _ in [0]*n:
    res += 'a'

これが4270msです。
Pypyでは文字列関連でこういうことがあるので提出前に絶対コードテストに通すべきです。
ちなみに上のをやりたい場合はjoin()を使うと早いです。
 
いくつかのライブラリ
凝ったことをするとREになることがあります。(累積和の第2引数など)
心配なら提出前にコードテストで試すべきです
 
deque
やや遅い
 
atcoder創世期
昔の昔の問題はメモリ制限が64MB(今は1024MB)なので、MLEになることがあるようです
 
リスト内包表記
Pythonだと圧倒的に高速化するリスト内包表記ですが、Pypyだと逆にやや遅くなります。

n = 10**7

a = [0]*n
for i in range(n):
    a[i] = i*2
n = 10**7

a = [i*2 for i in range(n)]

 
 べた書き(上):258 ms
 リスト内包表記(下) : 511ms

 
注意なのですが、Pythonで書く場合は下の方が圧倒的に早いです。
 

型をあべこべにするのはPypyだと実行速度に割と影響が出てしまいます。
よくあるのは、グラフの最短距離問題等の初期化としてfloat("inf")を用いますが、これは10**9とかの適切なint値にした方がいいです。


他にこれもあるよ!!って人はTwitter(じゅっぴー (@juppyjappy) | Twitter)でもコメントでもいいので教えてもらえると幸いです。
 
PypyでREになるの怖い...とかREって出たけどどこがダメなの?ってときは コードテストで試しましょう。
 
Pypy関係ないですが、サンプルは通るのにRE吐くときは
 ・存在しない要素のremove
 ・for文とかでのlistの範囲外
 ・無限ループ
 ・再帰限界
 ・n=1とかのそういうとき
 ・0で割ってる
 ・最初の条件分岐でa[1]とか使ってる
 メモリエラー
とかが多いですね。
 
 
 

Python定数倍高速化のテクニック

ここからが本番です。
 
まずこれを読んだことのない人は読みましょう。とても分かりやすくかつ有益です。
www.kumilog.net
入力が10^5個以上の時は必ずやるべきsys.stdin.readlineは、実は文字列の入力に対して"\n"まで含んでしまうので注意が必要です。

import sys
input = sys.stdin.readline

s = list(input()) #abc

print(s) #['a', 'b', 'c', '\n']

 
これ以外に知っておいた方がいいことを書いておきます。
 
 
if __name__ == '__main__':

def main():
    """"ここに今までのコード"""

if __name__ == '__main__':
    main()

とすると高速になります。
グローバル名前空間の変数アクセスの話です。
 
 
tupleを使う
ざっくり言うと「要素の変更ができない代わりに高速な配列」です。listより生成も速いです。
ダイクストラなどの「配列の出し入れ」をする問題など、使えるところは割とあります。
 
 
ライブラリをちゃんと使う
同じことをするのにもライブラリを使った方が速いです。
okumuraさんの以下の記事に一通り目を通しておくとよいです。
Pythonで競技プログラミング -ライブラリ編- - Qiita
 
scipyのダイクストラ等は以下の過去記事に書いてます。
(グラフ問題は基本Pypy一択ですが、、、)
scipyのFloyd-WarshallとDijkstraのすすめ Python 競技プログラミング Atcoder - じゅっぴーダイアリー
 
 
enumerate
for e in a: という書き方が速くなるとありましたが、これの補足としてindexと要素を知りたいときはenumerateを使うといいです。
 
小さいループは外側に
二重でfor文を回すときは小さいほうを外側にしましょう。
例えばn = 10**7で以下のループを回すと、

for i in range(n):
    for j in range(2):
        a += 1
for i in range(2):
    for j in range(n):
        a += 1 

 
 小さいループが内(上):5286 ms
 小さいループが外(下) : 1923ms

 
のような差がでます。

配列をできるだけ軽くする
軽くするというのは次元的な意味でも処理的な意味でもです。
基本的にPythonで定数倍の為にTLEするのは配列関連です。
 
よくある手法は「三次元の配列を二次元の配列何個かにする」というものです。
例えばDPで
 dp[i][j][k] = dp[i-1][j-1][k-1]*dp[i-1][j][[k]%mod
みたいな遷移があった時、これを3次元DPで実装するのではなく、下のように2次元配列二つで実装します。
 

"""dp = [[[0]*n for i in range(n)] for j in range(n)]"""
dp_f = [[0]*n for i in range(n)]
dp_s = [[0]*n for i in range(n)]

"""
for i in range(1,n):
    for j in range(1,n):
        for k in range(1,n):
            dp[i][j][k] = dp[i-1][j-1][k-1]*dp[i-1][j-1][k-1]%mod
"""   
for i in range(1,n):
    for j in range(1,n):
        for k in range(1,n):
            if i%2 == 0:
                dp_f[j][k] = dp_s[j-1][k-1]*dp_s[j][[k]%mod
            else:
                dp_s[j][k] = dp_f[j-1][k-1]*dp_s[j][k]%mod

(偶奇の場合分けの部分はもう少し早く書けますが、ここでの趣旨ではないのでわかりやすさ重視です)
 
. をなくす
よく見るのは append = res.append とローカルに格納していくものです。
例えばPythonでresという配列に 1 を10^7回appendするとき、

r_append = res.append

for i in range(n):
    r_append(1)

 

for i in range(n):
    res.append(1)

を比べると、
 r_append(上) :788ms
 res.append(下) :1049 ms

ほど変わります。
 
⑧listの出し入れをなくす
listの出し入れをする際(出し入れする配列自体は変更しない)、tupleを使うといいと上で書きましたが、
組を出し入れしない方法もあります
例えば、

n = 10**7
tank = []
for i in range(n):
    tank.append([65656,32131])
n = 10**7
tank = []
for i in range(n):
    tank.append((65656,32131))
n = 10**7
tank = []
for i in range(n):
    tank.append(65656*(10**8)+32131)

では、
上 : 5600ms(Python) , 2968ms(Pypy)
真ん中 : 1067ms(Python) , 1096ms(Pypy)
下 : 1077ms(Python) , 552ms(Pypy)
となっています。

 
配列の出し入れをする問題は、グラフで重みの小さな辺から見ていくダイクストラの応用みたいな問題でよく見ますが、
この上のは、TTPC2019で🦍さんとチームを組んだ際に指摘されて知りました。感謝ゴリ~
 

最後に

Pythonの定数倍高速化の闇は深いので、最終的には強い人のコードをみることで学んでいくしかないです。
自分も強い人から技術を盗んできました。ここでは最後に自分がコードをよく見せていただいている方をあげておきます。
いつも本当にありがとうございます。
 
prd_xxxさん:Pythonのlanguage Ownerです。500 streak(500ACじゃないよ!! streakだよ!!)というとてつもない記録を持ってます。🦍。
okumuraさん:個人的に一番コードがきれいだと思ってます。背伸びしていつもコードの書き方を真似させていただいています。
htkbさん : コードはもちろんTwitterなどで、Pythonの勉強において一番お世話になった方です。本当に何から何まで教えていただきました。
 

*1:経験則によるものやコミュニティで見かけたもので特にこれといったソースはないです。日本語か英語でPypyについての有用な文献あったら教えてほしいです。