Python における element in list が遅い件について

はじめに

以前の記事で、pythonにおいてelement in listという書き方は非常に遅いので気を付けましょうと書きました。これについて、実例とともに見ていきたいと思います。

element in listについて

pythonでは、リストにある要素があるかどうかを次のように判定することが可能です。

>>> 1 in [1, 2, 3] # True
>>> 4 in [1, 2, 3] # False

ただし、この処理は非常に遅いことが知られています（おわりに参照）。

例題

ABC129-C を扱います。まだ問題を解いたことがない方はネタバレ注意です。

問題概要

N段の階段があり、現在、0段目にいます。一度に移動できるのは、一段もしくは二段です。いま、a1, a2, ..., aM段目の床は壊れていて足を踏み入れることができません。壊れている床を踏まずにN段目までたどり着く方法は全部で何通りあるか求めてください。ただし、答えは10^9 + 7で割った余りを提出してください。問題へのURL

制約

• 1 <= N <= 10^5
• 0 <= M <= N - 1
• 1 <= a1 < a2 < ... < aM <= N - 1

解答

基本的な方針としては、dpテーブルを用意してあげて、dp[i]にi段目の階段にたどり着くための移動方法（を10^9 + 7で割った値）となるように構成してあげればよいと思います。このとき、床が壊れているような階段kに対して、dp[k] = 0となることに注意をしながら、dpテーブルを埋めていってあげます。

はじめに、TLEになってしまう書き方を提示します。

n, m = map(int, input().split())
a = [0] + [int(input()) for _ in range(m)]
MOD = 10 ** 9 + 7

dp = [0] * (n + 1)
dp[0] = 1

# 配列 a に 1 が含まれているかどうか
if 1 in a:
    dp[1] = 0
else:
    dp[1] = 1

# dp デーブル埋め
for i in range(2, n + 1):
    if i in a:
        continue
    dp[i] = dp[i - 1] + dp[i - 2]
    dp[i] %= MOD

print(dp[-1]) # dp[n] と同じ

このコードをPython3で提出するとTLEとなってしまいます。これは、dpテーブルを埋める際に、毎回の反復でiが配列aに含まれているかどうかの判定を行っていることが原因です。

---

そこで、次のように書き方を変更してあげましょう。すなわち、毎回の反復でiが配列aに含まれているかどうか判定するのではなく、事前に床iが壊れているかどうかのboolean配列を用意してあげて、それを利用しましょう。

n, m = map(int, input().split())
a = [0] + [int(input()) for _ in range(m)]
MOD = 10 ** 9 + 7

ok = [True] * (n + 1) # 床 i が壊れている ⇒ False
for v in a: # もちろん入力の段階で処理をしてもよい
    ok[v] = False

dp = [0] * (n + 1)
dp[0] = 1
# 配列 a に 1 が含まれているかどうか（ここで一度 in を使うくらいはいいでしょう）
if 1 in a:
    dp[1] = 0
else:
    dp[1] = 1

# dp テーブル埋め
for i in range(2, n + 1):
    if ok[i] == True: # 変更箇所
        dp[i] = dp[i - 1] + dp[i - 2]
        dp[i] %= MOD

print(dp[-1]) # dp[n] と同じ

これをPython3で提出してあげると、200msくらいでACとなります。

おわりに

おわりに、element in listという書き方がどれくらい遅いのか、jupyter notebookで確認しておきましょう（思い付き）。

次のようなコードで実験してみましょう。lstという配列の要素に-1があるかどうかを判定する操作を1000回行うコードです。 なお、lstは事前に

lst = list(range(10 ** 5))

のように定義しておきます。

%%timeit
for i in range(1000):
    if -1 in lst:
        print('')

結果は次のようになりました。

1.67 s ± 211 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

平均で1.67 sかかっています。今回は反復回数を1000としましたが、上で扱った問題では最高で10 ** 5回程度の反復を要します（Nの制約条件より）。TLEになるのも納得できます。