PyPyを速くする（3） ジェネレータ

例えば素因数分解を与えて約数を全て列挙するコードを書きます。こんな感じでしょうか。

def pows(n, e):
    yield 1
    q = 1
    for _ in xrange(e):
        q *= n
        yield q

def divisors(fs):
    if not fs:
        yield 1
    else:
        p, e = fs[0]
        qs = pows(p, e)
        for n in divisors(fs[1:]):
            for q in qs:
                yield n * q

10^6までの約数の個数を数えると、Pythonで3.4秒、PyPyで2.1秒でした。PyPyはジェネレータが遅いので、リストにしたほうが速くなります。

def pows2(n, e):
    qs = [ 1 ]
    q = 1
    for _ in xrange(e):
        q *= n
        qs.append(q)
    return qs

def divisors2(fs):
    if not fs:
        return [ 1 ]
    else:
        p, e = fs[0]
        qs = list(pows2(p, e))
        return [ n * q for n in divisors(fs[1:]) for q in qs ]

Pythonで4.5秒、PyPyで1.6秒になりました。
まとめると、

	Python	PyPy
ジェネレータ	3.4秒	2.1秒
リスト	4.5秒	1.6秒

【結論】
ジェネレータを使うと美しく書けます。また、長い列を生成して途中で切りたいときもあるでしょう。しかし、上の例のように短くて全部なめることもよくあります。こういうときはスピードアップのために、ジェネレータをやめてリストにすることも考えましょう。