Zenn
😇

【Python】競プロにおける配列の初期化の注意点

2023/02/16に公開
Python
AtCoder
tech

float('inf')やmath.infを使うな

タイトルの通りで、dpなどの配列の初期化時にfloat('inf')などを使って初期化すると、若干遅くなります。制約がそんなに厳しくないような問題だと、そこまで気にする必要はないのですが、想定解がO(10^8)とかで、特にPythonでは厳しい制約の場合には、float('inf')などで初期化すると、それがボトルネックとなってTLEしてしまうことがあります。
そこで、float('inf')、math.inf、INF=pow(10, 12)で配列を初期化する際の計算時間をAtcoderのコードテストを用いて比較してみました。

テストコード

float('inf')で初期化時の検証

N = ??
dp = [float('inf') for _ in range(N)]

math.infで初期化時の検証

import math
N = ??
dp = [math.inf for _ in range(N)]

INF=pow(10, 12)で初期化時の検証

N = ??
INF = pow(10, 12)
dp = [INF for _ in range(N)]

いずれも5回実行した際の平均値を検証結果としています。

Python(3.8.2)での比較

xは10s以上を示す。

N float('inf') math.inf INF=pow(10, 12)
10 16.6ms 17.4ms 17.4ms
10^2 17.8ms 18.4ms 18.8ms
10^3 19.4ms 18.6ms 18.4ms
10^4 23.6ms 19.6ms 18.0ms
10^5 41.2ms 32.0ms 23.8ms
10^6 175.2ms 84.6ms 63.4ms
10^7 1535.8ms 602.2ms 412.0ms
10^8 x 5809.8ms 3888.0ms

PyPy3(7.3.0)での比較

N float('inf') math.inf INF=pow(10, 12)
10 68ms 114.4ms 63.6ms
10^2 64.6ms 63.6ms 63.8ms
10^3 71.2ms 69.0ms 62.2ms
10^4 71.0ms 135.4ms 69.2ms
10^5 83.0ms 73.6ms 68.8ms
10^6 160.2ms 111.6ms 110.8ms
10^7 955.6ms 555.6ms 541.8ms
10^8 x 5380.0ms 5512.0ms

結論

math.infはまだしも、 「float('inf')は使うな」 ということですね。INF=pow(10, 12)としておくのが無難だと思います。とはいえ、これが直に効いてくる問題は滅多にないとは思いますが。
float('inf')だとTLEで、INF=pow(10, 12)だとACとなる問題は例えばABC 192 F-Potionですかね。これは1 \leq N \leq 10^2で、空間複雑度O(N^3)の配列dpをN回作成して解く、O(N^4)が想定解となる問題です。
自分が提出したコードを載せます。コードの詳細は割愛しますが、これらのコードの違いは、配列の初期化の方法のみです。INF=pow(10, 12)で配列を初期化した場合はAC、float('inf')で初期化した場合はTLEになってしまうのが分かります。

ACコード
TLEコード

Discussion