itertoolsで素数計算

Python 30分プログラム

30分プログラム、その651。itertoolsで素数計算をしてみた。

今日、ちょっと動揺することがあったので、心を落ち着けるために素数を数えるコードを書いてみました。ネタのつもりだったんですけど、わりと落ち着きました。よかったです。

で、その辺のことはおいといて、Pythonの話をしましょう。

何度も書いてるけど、Pythonのストリームと遅延リストは似て非なるものです。何も考えずエラトステネスをふるいを書くと、関数呼び出しが止まりません。

def sieve(xs):
    x = xs.next()
    ys = sieve(ifilter(lambda y: y % x != 0,xs))
    return chain([x],ys)

この問題は、明示的にyieldを呼んでやればきっと回避できますが、それだと30分もかからず解けてしまいます。それにあんまりスマートじゃない気もします。
というわけで、遅延関数適用iapplyを作ってみました。

使い方

primes = sieve(count(2))
print take(10,primes)
# [2, 3, 5, 7, 11, 13, 17, 19, 23, 29]

ソースコード

#! /usr/bin/python
# -*- mode:python; coding:utf-8 -*-
#
# prime.py -
#
# Copyright(C) 2009 by mzp
# Author: MIZUNO Hiroki / mzpppp at gmail dot com
# http://howdyworld.org
#
# Timestamp: 2009/09/01 21:37:55
#
# This program is free software; you can redistribute it and/or
# modify it under MIT Lincence.
#

from itertools import *

class iapply():
    def __init__(self,f,*args):
        self.f = f
        self.args = args
        self.ret  = None

    def __iter__(self):
        return self

    def next(self):
        if self.ret == None:
            self.ret = apply(self.f,self.args)
        return self.ret.next()

def sieve(xs):
    x = xs.next()
    ys = iapply(sieve,ifilter(lambda y: y % x != 0,xs))
    return chain([x],ys)

def take(n, seq):
    return list(islice(seq, n))

primes = sieve(count(2))
print take(10,primes)

参考