【競プロ用】Python高速化Tips

最近Pythonで競プロの勉強をしている中で（TLEを叩きまくりながら）気付いた遅い処理とそれを**速くする方法（ACした方法）**を備忘録代わりに書きます。緑コーダーなので難しいことはわかっていませんがご容赦下さい。なお、実行時間はAtCoderのコードテストでPython3.8.2を使い測定しました。

書いていないこと

• 言語（Python）とは関係ないこと（アルゴリズム）
  • 二分探索を使う
  • メモ化する　など
• 以前から知っていたこと、まだ実感を得られていないこと
  • Python高速化の方法は色々な方が発信されていると思いますので、興味のある方は「Python 競プロ 高速化」あたりでググってみて下さい。自分は特にPython で AtCoder をするあれこれを参考にさせて頂きました

Tips

PythonではなくPyPyを使う

PythonではなくPyPyを使った方が基本的には速くなります。【競プロ】PythonとPyPyの速度比較が参考になりましたので是非見てみて下さい。AtCoderではPyPyが使えるため$O(10^6)$くらいの計算量になるときはPyPyを使うことが多いです。ただし、2021年9月時点ではNumPyとの共存はできないようです（numpy.ndarrayも便利で好きなため残念）。また、PyPyは再帰関数は遅いのでその点も注意です。AtCoder Beginner Contest 149 F - Surrounded NodesではPyPyではTLEでしたがPythonではACしました。

2022/2/26追記

PyPyは再帰関数は遅いのでその点も注意です。

と書いたのですが、どうも

import pypyjit
pypyjit.set_param('max_unroll_recursion=-1')

と書くとPyPyの再帰関数が速くなるようです。仕組みはわかりませんが、ループ展開というのが関係してそうです。↓見るとメモリ使用量も少なくなっていますね。

存在チェックをしたいときlistではなくsetを使う

setは重複を許しませんし、二分探索も可能なのでlistよりも高速に存在チェックできます。"X in list"(線形探索・二分探索)と"X in set"の計算時間について調べてみたに詳しく書かれていますので見てみて下さい。さらに言えば、存在するかどうかをリストA[a]で管理すれば$O(1)$で判別可能です（その分メモリは使います）。

# 遅い -> 554 ms
A = []
for _ in range(10**2):
    for a in range(10**3):
        if a not in A:
            A.append(a)

# 速い -> 34 ms
A = set()
for _ in range(10**2):
    for a in range(10**3):
        if a not in A:
            A.add(a)

値の判定の際にrangeではなく大なり小なりを使う

当たり前のことですが、結構速度に差がありますので大なり小なりを使いましょう。

# グリッドの問題でよく出てくる書き方
H, W = 1000, 1000

# 遅い -> 590 ms
for h in range(H):
    for w in range(W):
        if not (h in range(H)) or not (w in range(W)):
            continue

# 速い -> 147 ms
for h in range(H):
    for w in range(W):
        if not (0 <= h < H) or not (0 <= w < W):
            continue

listのコピーにcopy.deepcopy()ではなくスライスを使う

これも普段から後者を使うのがいいと思います。

import copy
old_list = [x for x in range(10**3)]

# 遅い -> 377 ms
for _ in range(10**3):
    new_list = copy.deepcopy(old_list)

# 速い -> 21 ms
for _ in range(10**3):
    new_list = old_list[:]

2値の比較の際にはmin, maxの代わりにifを使う

つい便利なので前者を使いがちですが、後者の方が早いです。

# DPの問題でよく出てくる書き方
dp = [0] * 10**3

# 遅い -> 246 ms
for x in range(10**3):
    for i in range(10**3):
        dp[i] = max(dp[i], x)

# 速い -> 140 ms
for x in range(10**3):
    for i in range(10**3):
        if x > dp[i]:
            dp[i] = x

insertする際にはlistではなくarrayを使う

arrayは普段は使いませんが、TLEになったらこの手も考えます。ちなみに単なる探索の場合（insort_leftの代わりにbisect_leftとした場合）、前者は49 ms、後者は59 msと逆転しましたので、探索だけの場合はlistを使うのがよさそうです。

import bisect, array

# 遅い（というか普通はこう書く） -> 451 ms
q = list(range(0, 10**5, 2))
for i in range(1, 10**5, 2):
    bisect.insort_left(q, i)

# 速い -> 142 ms
q = array.array('i', list(range(0, 10**5, 2)))
for i in range(1, 10**5, 2):
    bisect.insort_left(q, i)

ちなみにAtCoder Beginner Contest 217 D - Cutting Woodsではlistを1つ使い挿入、探索する方法ではTLEしたのですが、listを2つに分割する（小さな数字を管理するlistと大きな数字を管理するlistに分ける）方法でACしました。また、コンテスト後に試しましたがarrayを使う方法でもACしました。

【2021/9/16追加】2次元、3次元リストではなく1次元リストを使う

複雑になりバグを生みやすくなるので積極的に使いたくはないですが、TLEになったらこの手も考えます。特に3次元リストになるとリストの初期化に時間がかかるようです。反対にループの中でリストを利用する際、1次元リストではcost[(h*W+w)*4+i]といった計算が必要になるため、3次元リストを使うよりも遅くなります。043 - Maze Challenge with Lack of Sleep（★4）では3次元リストを1次元リストにすることでAC取れました。

import math
H = 1000
W = 1000
INF = math.inf

# 遅い（というか普通はこう書く） -> 1215 ms（初期化だけで566 msかかる）
cost = [[[INF] * 4 for _ in range(W)] for _ in range(H)]
for h in range(H):
    for w in range(W):
        for i in range(4):
            if cost[h][w][i] == INF:
                pass

# ちょっと速い -> 1007 ms（初期化は67 msだが、その後利用する際に時間がかかる）
cost = [INF] * (4 * W * H)
for h in range(H):
    for w in range(W):
        for i in range(4):
            if cost[(h*W+w)*4+i] == INF:
                pass

【2021/9/16追加】PyPyを使う場合、main関数を使わないようにする

Pythonの場合、main関数を使った方が速いと聞いたことがあるのですが、試してみたところPyPyの場合はmain関数を使わない方が速かったです。理由はわかりません・・・。055 - Select 5（★2）をPyPyで解く際、main関数を外すことでACできました。このケースで言えば5重ループを書いた方がより速くなります。

# 遅い（main関数を使う） -> PyPyでは265 ms（Pythonでは431 ms）
def main():
    from itertools import combinations

    P, Q = 100, 99
    A = range(50)
    ans = 0
    for a, b, c, d, e in combinations(A, 5):
        if a % P * b % P * c % P * d % P * e % P == Q:
            ans += 1
    print(ans)

if __name__ == '__main__':
    main()

# 速い（main関数を使わない） -> PyPyでは159 ms（Pythonでは728 ms）
from itertools import combinations

P, Q = 100, 99
A = range(50)
ans = 0
for a, b, c, d, e in combinations(A, 5):
    if a % P * b % P * c % P * d % P * e % P == Q:
        ans += 1
print(ans)

【2021/9/8追加】不要ならばmath.infやfloat('inf')を使わず-1などを使う（BFSの初期値など）

math.infと-1の差でTLE -> ACということはよっぽどないと思いますが、-1を使う方が多少速くなります。

import math
N = 10**6

# ちょっと遅い -> 109 ms
for i in range(N):
    if i == math.inf:
        pass

# ちょっと速い -> 76 ms
for i in range(N):
    if i == -1:
        pass

# 中間 -> 93 ms
INF = math.inf
for i in range(N):
    if i == INF:
        pass

【2021/9/8追加】ループ内の処理はシンプルに保つ（if文での判定は行わずループ後にうまく処理するなど）

言語（Python）とは関係ないことですし当たり前のことなのですが、問題解いている最中には気付かなかったので・・・。具体的には2010年 日本情報オリンピック春合宿OJ finalsでこれを使いACしました。

# ループの中でグループ数を数え判定しようとしたところTLEになった
for c, a, b in G:
    if uft.group_count() <= K:
        break
    if not uft.same(a, b):
        uft.union(a, b)
        ans += c

# ループ後に必要な数値のみ切り出し答えとしたところACした
for c, a, b in G:
    if not uft.same(a, b):
        uft.union(a, b)
        cost.append(c)
ans = sum(cost[:N-K])

【番外編：メモリ超過を避けるための方法】

AtCoderでMLEになることは少ないですが、昔の問題を解いているとたまに起こります。もしも発生すれば以下の案を検討するのがいいと思います。

• 隣接行列ではなく隣接リストを使う
• 保持する必要のない情報は保持しない（例えばtupleを使い2つの情報を保持していてMLEになれば、1つの情報だけで成立しないかを考える）