2012年9月4日火曜日

【Python】One-liner でフィボナッチ数列を返すジェネレータを作る

今日学校で議論していたのがこの話題。
Python にはジェネレータというものがあります。そこで、フィボナッチ数列を1つずつ返すジェネレータを作りたいという話になりました。

まずジェネレータを簡単に紹介します。
(i * 2 for i in range(0, 10))
とすると、これは 0, 2, ..., 18 を返すジェネレータになります。
ジェネレータはfor文で便利に扱えます。
g = (i * 2 for i in range(0, 10))
for x in g:
        print x
とすると、出力は
0
2
4
6
8
10
12
14
16
18
となります。

さて、ジェネレータを関数形式で書くことも可能です。
フィボナッチ数列を1つずつ返すジェネレータは次のように書けます。
def fib():
        a, b = 0, 1
        while True:
                yield a
                a, b = b, a + b

g = fib()
for i, f in enumerate(g):
        if i == 10:
                break
        print f
enumerateというのは、シーケンスの全要素に対し、各要素の位置をくっつけたタプルを返してくれる関数です。ここでは第10項までのフィボナッチ数列を求めるために使っています。

これ、 One-liner (一行プログラム)で書けたら素敵ですよね。という話になったのです。そのための構文があってもいいもんだなあ、と。
例えばこんな構文があったらどうかという結論になりました。
gen init ret next
genはキーワード、initは初期値、retは現在の値から yield するための値を取り出す関数、nextは現在の値から次の値を求める関数です。
フィボナッチ数列だったら次のように書けます。もちろん、今の Python にはこんな構文無いのでエラーですが。
g = gen (0, 1) (lambda (a, b) : a) (lambda (a, b) : (b, a + b))
for i, x in enumerate(g):
        if i == 10:
                break
        print x
これを実行すると次のような出力になることが期待されます。
0
0
1
1
2
3
5
8
13
21
34

素敵ですね。
ということで、無いなら作ってみました。
def gen(init, ret, next):
        while True:
                yield ret(init)
                init = next(init)

g = gen((0, 1), lambda (a, b) : a, lambda (a, b) : (b, a + b))
for i, x in enumerate(g):
        if i == 10:
                break
        print x
gen関数の定義は、gen構文が無いために仕方なく作った関数なので、そこを指摘して「One-liner じゃないじゃん!」と突っ込むのはやめてください。
ミソはgen関数を用いてジェネレータを生成しているところ。
while True:とかyieldとか書かなくても、結構複雑なジェネレータを1行で書くことができました。

おまけ:One-liner で素数ジェネレータ
https://gist.github.com/3618684

0 件のコメント:

コメントを投稿