結局pythonのforループはどのように書けば早いのか

結局pythonのforループはどのように書けばいいのか悩むことが多く、個人的に少し試してみましたので備忘録として書いておきたいと思います。

forループを高速にするための記法としては、「リスト内包表記」が有名ですが、ジェネレータを使った書き方も試してみます。

三重forループを以下4つの書き方を試してみます

比較してみる

N = 300
lst=[]
for X in range(N):
  for Y in range(N):
    for Z in range(N):
      lst.append(X+Y+Z)
ans = sum(lst)

forループを3つ重ねてみました。
5回実行したときの平均所要時間は6.06[sec]でした。

先ほどのコードをリスト内包表記で書いてみます。

N = 300
tmp_list = [X+Y+Z for X in range(N) for Y in range(N) for Z in range(N)]
ans = sum(tmp_list)

5回実行したときの平均所要時間は4.36[sec]でした。

リストを作る必要はないのでジェネレータ式に変更してみます。リストの[]を()に変更するだけです。

N = 300
tmp_list = (X+Y+Z for X in range(N) for Y in range(N) for Z in range(N))
ans = sum(tmp_list)

5回実行したときの平均所要時間は3.54[sec]でした。
ジェネレータ式にすると、早くなりますね。

内包表記を用いることで記述を簡潔にできますが、処理によっては可読性を損なうことがあります。できればforループを使いつつ、高速化する方法が欲しいところです。
そこで下記のようにジェネレータを定義する関数を用意してみました。

  def sample():
    for X in range(N):
      for Y in range(N):
        for Z in range(N):
          yield X+Y+Z

  ans = sum(sample())

5回実行したときの平均所要時間は3.71[sec]でした。

遅くなるかと思いましたが、ジェネレータ内包表記とほとんど差がない実行速度にできました。可読性が悪くなるような内容ならこのようにヘルパー関数のようなものでジェネレータを作ってみるとよさそうです。

forループの書き方	実行速度(5回の平均)	メリット	デメリット
愚直にforループを回す	6.06	可読性がよい	遅い
リスト内包表記	4.36	やや早い	可読性を損なう場合がある
ジェネレータ内包表記	3.54	早い。メモリを節約できる。	可読性を損なう場合がある。ジェネレータのため複数回呼び出せない
愚直にforループを回しつつyieldで返す	3.71	早い。メモリを節約できる。可読性を担保できる。	ジェネレータのため複数回呼び出せない。

結論としては、
forループを書く際の基本方針は、
ジェネレータを用いて書く。基本は内包表記で書き、ただし可読性を担保できないようであれば、ジェネレータを関数で定義する。

もちろん、例外はあると思います。

以上です。

@intermezzo-frさんの記事 Pythonのリスト内包表記の速度を拝見すると、愚直なforループが遅いことの原因はappendの実行速度にあるそうです。そういう意味では、愚直にforを書いてyieldで返しても内包表記と実行速度がほとんど変わらないのは自然かもしれません。