自分への戒めも込めて
Pythonはとても便利。
でも下手な書き方すると計算が全然終わらないことがある。
一方でちょっとした書き方の工夫をするだけで計算が全然速くなったりする。
ここでは自分がかつて陥った「内包表記の計算が全然終わらない」ケースとその対処法をいくつか紹介する。
内包表記ってどんなだっけ？
例えば以下のようにlistを作成したいとする。
愚直にやるとこんな感じ
l = []
for i in range(10):
l.append(i)
print(l)
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]
これを内包表記を使うと以下のようにスッキリ書ける。
l = [i for i in range(10)]
Pythonはその言語の性質上ネストが深くなるとインデントもどんどん深くなりがち。
内包表記を使うとインデントを下げることなくすっきり書けるので便利。
JavaやCなどの言語からPythonに入門するとついつい最初みたいな書き方をしてしまうが、
内包表記を積極的に使っていこう。
時間経過の計測
今回は計算時間を測定するために
import time
start = time.time()
# 何らかの時間のかかる処理
t = time.time() - start
print(t)
という書き方をした。
毎回書くのは面倒なので省略して書いていることもあるが、察して欲しい。
罠1: listから探索するな
整数列のリストと偶数のリストを以下のようにそれぞれnum_list, even_listとする。
ここでnum_listからeven_listに含まれない要素だけを取り出して、奇数のリスト odd_listを作りたい。
以下のように書けそうだ。
num_list = [i for i in range(10)]
even_list = [2 * i for i in range(10)]
odd_list = [i for i in num_list if i not in even_list]
print(odd_list)
[1, 3, 5, 7, 9]
が、この書き方はeven_listが長くなると簡単に計算時間が爆発する。
以下のようにちょっと数字をいじっただけで、手元の環境だと40秒以上かかった。
num_list = [i for i in range(10000)]
even_list = [2 * i for i in range(10 ** 6)]
odd_list = [i for i in num_list if i not in even_list]
# time: 44.69495511054993
理由は簡単で、Pythonのlistは先頭から順に走査するためだ。
そのため二つのリストの長さをそれぞれM, NとするとO(MN)で時間がかかってしまう。
解決策は簡単で、探索対象のeven_listをset（集合）にしてしまえば良い。
setにすることで順番の情報は失われるが、今回はnum_listの各要素がeven_listに含まれるかどうかだけ
分かればいいので問題ない。
以下のように書き換えるだけで0.1秒で回すことができた。
num_list = [i for i in range(10000)]
even_list = [2 * i for i in range(10 ** 6)]
even_set = set(even_list)
odd_list = [i for i in num_list if i not in even_set]
# time: 0.10969114303588867
setやdictなどは内部にhashを持っているのでO(1)で回すことができる。
そのため上記のように書けばO(M)だけで回すことができるようになる。
内包表記に限った話ではないが、listのまま計算すると無駄に計算時間がかかることが多い。
順番を保持する必要がない場合は積極的にsetを使っていきたい。
罠2: loopの中で計算するな
さて、先程のsetを使った内包表記だが、少し気を抜くと以下のように書いてしまうかもしれない。
setの計算をいちいち変数に格納せずに内包表記の中でset(even_list)してみた。
num_list = [i for i in range(10000)]
even_list = [2 * i for i in range(10 ** 6)]
odd_list = [i for i in num_list if i not in set(even_list)]
# time: 255.01793599128723
しかし、この書き方はめちゃくちゃ時間がかかってしまっている。
なぜ？
これはloopのたびにlist-> setの変換処理が走ってしまっているためだ。
内包表記のままだと気付きづらいが、これは普通にfor文を書くと以下と同じだ。
num_list = [i for i in range(10000)]
even_list = [2 * i for i in range(10 ** 6)]
for i in num_list:
if i not in set(even_list):
l.append(i)
こう書けばnum_listのloopの回数だけset()の計算が行われていることは明らかだろう。
内包表記で書くとぱっと見set(even_list)を一度しか計算していないようにも見える（僕だけ？）が、
基本的にloopの外でできる計算は外でやるべきだ。
罠3: 長すぎるlistを使うな
num_listが長い場合、計算結果の最後にsliceで先頭10件を取ってきたいとする。
num_list = [i for i in range(10 ** 7)]
l = [2*i for i in num_list][:10]
# time: 1.055954933166504
[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
これも書き方を工夫できる。
10**7のlistを全部走査してから10件を切り取るよりも、最初からlistの先頭10件を切り取っておいた方が速い。
num_list = [i for i in range(10 ** 7)]
# method2
l = [2*i for i in num_list[:10]]
# time: 0.34849095344543457
細かい違いだが、num_listが長くなるほど、また2*iの計算部分が複雑になるほどこの効果は大きくなる。
この例だと最初に10件にsliceすれば出力も10件なので問題ないが、最初の奇数の列を作る場合はどうしよう？
つまり以下のように内包表記の中にif文なんかが入り込むと件数が変わってくるということだ。
odd_list = [i for i in num_list[:10] if i not in even_list] # 出力 5件
odd_list = [i for i in num_list if i not in even_list][:10] # 出力 10件
この例だとif文の処理が簡単なので件数は綺麗に半分だが、一般にはif文を当てはめて何件出力されるかは不明だ。
この場合は下の書き方で妥協するか？
いや、この場合でも内包表記内のnum_listをできるだけ短くすることで時間を短くできる。
以下のようなイメージだ。
num_list = [i for i in range(10 ** 7)]
small_num_list = num_list[:1000]
l = [i for i in small_num_list if i not in unknown_list][:10]
if文でどれくらいの要素が当てはなるかはわからないが、明らかに10**7個もlistが必要でないことはある。
そんなときはいったん短めのリストを作っておきそこからloopを回すことで時間を節約できる。
まとめ
今回紹介したのはPythonの中でもかなり初歩的な処理だ。
他にも
「listじゃなくてnp.arrayを使うと速い」
「CPUじゃなくてGPUを使うと速い」
「○○の処理は××のライブラリを使った方が速い」
など高速化する方法は無数に存在するわけだが、ちょっとした書き方の工夫も意外に馬鹿にならないので意識していきたい。
（自分は修行が足りてないので、気にしないとついついやってしまうことが多い）
「この計算はO(NM)だな？」とか「これはO(log(N))で回せそう」とか、特にloopを実装するときに計算量を意識する癖をつけると良いのかなーと思う。
おわり。