素数を列挙するイテレータ

単純に2から割り算を繰り返すものでもいいのですが、そこそこ効率がいいものにするため、最初にエラトステネスの篩をかけるバージョンを作ってみます。

まず2、3、5の倍数を削除した数列を1周期分(0〜29)だけ作成します。あとは周期30(=2×3×5)を加算しながら列挙していくイテレータを使えば、エラトステネスの篩を無限列に対して行ったのと同じことになります。

import operator
def make_sieve(prime):
    n = reduce(operator.mul, primes)
    L = range(n)
    for p in prime:
        L[::p] = [None]*len(L[::p])
    return n, filter(None, L)
def gen_sieved(n,L):
    base = 0
    while 1:
        for x in L:
            yield base+x
        base += n

関数make_sieve([2,3,5])は周期30と増分リスト[1, 7, 11, 13, 17, 19, 23, 29]を返します。gen_sieved()はそれを受け取り2、3、5の倍数を落とした無限数列(1、7、11、13、17、19、23、29、31、37、41、43、47、49、53、59、…)を返すイテレータを作ります。

このgen_sieved()を被除数ループと除数ループの両方に使います。最初に1を返すのをnext()で飛ばしてやる必要があります。

もし割り切れたときは除数ループを中断して次の被除数に移ります。割り切れないで除数が被除数の平方根より大きくなったらこれ以上は割り切れることはありえません。被除数を素数としてyieldで返してから次の被除数に移ります。

def gen_prime(prime=[2,3,5]):
    for p in prime:
        yield p
    n,L = make_sieve(prime)
    dividend = gen_sieved(n,L)
    dividend.next()
    for dd in dividend:
        divider = gen_sieved(n,L)
        divider.next()
        for dr in divider:
            if dr*dr > dd:
                yield dd
                break
            if dd % dr == 0:
                break

gen_prime()は無限に素数を返し続けるので、実際に使うにはislice()などを使って個数を制限することになります。
たとえば次のようにすれば10000番目の素数を表示します。

from itertools import islice
print islice(gen_prime(), 9999, None).next()
104729

ふるい落としはデフォルトで[2,3,5]にしましたが、[2,3,5,7]とすれば7の倍数も、[2,3,5,7,11]ならさらに11の倍数までもふるい落としてから計算をはじめます。ただしあまりやりすぎると最初に巨大なリストを操作するためかえって遅くなってしまうかもしれません。