Pythonのジェネレータ式/filter関数でリストを効率よく検索する

概要

例えば、整数が格納されているリストから偶数のみ取り出して、先頭の3件のみ取り出す処理を考えてみます。

なお、偶数を取り出す（判断する）関数は共通で以下を利用します。

def is_even(x):
    print("in is_even({x})".format(x=x))
    return x % 2 == 0

リスト内包表記

まず、一番オーソドックスなリスト内包表記パターンです。

print("# リスト内包表記")
print("Prepare")
l = [x for x in range(10) if(is_even(x))]
print("Prepared")

# 先頭3件を取得
assert (l[0], l[1], l[2]) == (0, 2, 4)

特に難しい箇所はありません。
このコードを実行すると、以下のように出力されます。

# リスト内包表記
Prepare
in is_even(0)
in is_even(1)
in is_even(2)
in is_even(3)
in is_even(4)
in is_even(5)
in is_even(6)
in is_even(7)
in is_even(8)
in is_even(9)
Prepared

リスト内包表記が実行された時点で、is_even関数が全ての要素分実行されていることが分かります。
これはつまり、もしis_even関数が1回1秒掛かるような処理の場合、必ず10秒掛かるということになります。

ジェネレータ式

上記のリスト内包表記の処理を、ジェネレータ式に書き換えてみます。

print("# ジェネレータ式")
print("Prepare")
gen_exp = (x for x in range(10) if(is_even(x)))
assert type(gen_exp) == GeneratorType
print("Prepared")

# 先頭3件を取得
assert (next(gen_exp), next(gen_exp), next(gen_exp)) == (0, 2, 4)

実行結果は以下のようになります。

# ジェネレータ式
Prepare
Prepared
in is_even(0)
in is_even(1)
in is_even(2)
in is_even(3)
in is_even(4)

さて、出力結果がリスト内包表記の場合と全く異なりますね。
まず目につくのが**PrepareとPreparedの間に何も出力されていない**という点です。
これはつまり、ジェネレータ式を以下のように記述した場合、まだis_evenは実行されていないということになります。

gen_exp = (x for x in range(10) if(is_even(x)))

では、いつis_evenが実行されるとかというと、それは生成されたジェネレータ式に対してnextが呼ばれたタイミングです。
ジェネレータ式で指定しているis_even関数が、Trueを返すまで、リスト（などのコレクション）の要素を順番にis_evenに渡してくれます。もしTrueが返って来たらその時点でその際のリストの値が返されます
つまり、

next(gen_exp)

が初めて呼ばれた場合、range(10)の最初の要素0がis_evenに渡されます。
そして0はTrueになるので単純にこの値がジェネレータ式から返されます。

2回目に

next(gen_exp)

が呼ばれた場合はその次の要素から処理が開始されます。
ジェネレータ式が前回Trueが返された場所でループを止めてくれているイメージです。そこから処理が再開されます。
つまり、2番めの要素である1がis_evenに渡されます。しかしこれはFalseを返すので、さらに次の要素である2がis_evenに渡されます。
これはTrueなので、この2が結果としてジェネレータ式からreturnされます。
この時点で、is_evenは要素0、1、2のために合計3回呼ばれています。

そして再度

next(gen_exp)

を呼ぶと、前回同様その続きから処理が再開されるので、次の要素3がis_evenに渡され、Falseになるのでさらに次の要素4がis_evenに渡されて、これはTrueなので、結果としてこの4がジェネレータ式からreturnされます。
今回は要素3と4の合計2回is_evenが呼ばれたので、今のところ合計5回is_evenが実行されています。

さて、すでに3回next(gen_exp)を実行してその結果を得ているので、整数が格納されているリストから偶数のみ取り出して、先頭の3件のみ取り出す、という当初の目的を達成できてしまいました。
今回、is_evenが呼び出された回数は合計5回なので、リスト内包表記の場合同様1回1秒と仮定すると、処理時間はたったの5秒となります。
このような、全ての要素が必要となる訳ではない、というような処理の場合にジェネレータ式が重宝するのでは、と思います。

filter関数

基本的にジェネレータ式と同じです。

print("# filter関数")
print("Prepare")
filterd = filter(is_even, range(10))
assert type(filterd) == filter
print("Prepared")

# 先頭3件を取得
assert (next(filterd), next(filterd), next(filterd)) == (0, 2, 4)

実行結果も以下のようになり、ジェネレータ式同様の動作になっていることが解ると思います。

# filter関数
Prepare
Prepared
in is_even(0)
in is_even(1)
in is_even(2)
in is_even(3)
in is_even(4)

おまけ（nextのメリット）

さて、ジェネレータ式でもfilter関数でも、nextを実行しました。
もし値がもう無いよ、という状態でさらにnextを実行すると、デフォルトではStopIterationというエラーが発生してその時点でプログラムが止まってしまいます。

# ジェネレータ式のやつの結果を以下のようにすると。。。
assert (next(gen_exp), next(gen_exp), next(gen_exp)) == (0, 2, 4)
assert (next(gen_exp), next(gen_exp), next(gen_exp)) == (6, 8)

以下のようなエラーになります。

Traceback (most recent call last):
  File "ites_test.py", line 36, in <module>
    assert (next(gen_exp), next(gen_exp), next(gen_exp)) == (6, 8)
StopIteration

実はこのnext関数、第2引数には値が無かった際に返す値を指定することが出来ます。
そうするとエラーを抑制できる上に、値がない場合に利用するデフォルトの値までセット出来て一石二鳥です。
以下のように記述します。

assert (next(gen_exp), next(gen_exp), next(gen_exp)) == (0, 2, 4)
assert (next(gen_exp, None), next(gen_exp, None), next(gen_exp, None)) == (6, 8, None)

おまけ（ジェネレータをforで使う）

当然ジェネレータ式もfor文で使えます。

gen_exp = (x for x in range(10) if(is_even(x)))
result = []
for x in gen_exp:
    result.append(x)
assert result == [0, 2, 4, 6, 8]

この場合、リストの最後まで処理が走るので、is_evenの実行回数はリスト内包表記と同様になります。

まとめ

今まで余り意識することのなかった無限リスト等がこういう場面（ジェネレータ式）で役に立つんだな！と初めて実感出来ました。
Python素敵です！