自動的に拡張される素数リストの実装

Pythonの素数ジェネレータの実装は結構有名なはず。ProjectEulerのために、ランダムアクセス可能で必要に応じて拡張される素数リストを、素数ジェネレータを利用して作成した。ジェネレータの実装は4 TopCoderより。

from itertools import ifilter
from itertools import count

class Prime(object):
    def __init__(self):
        self.gen = sieve()
        self.primes = []
    def __getitem__(self,index):
        if len(self.primes) <= index:
            shortage = index - len(self.primes) + 1
            new_primes = [p for i,p in zip(range(shortage),self.gen)]
            self.primes.extend(new_primes)
        return self.primes[index]

def sieve():
    g = count(2)
    while True:
        prime = g.next()
        yield prime
        g = ifilter(lambda x,prime=prime: x%prime,g)

使い方はこんな感じ。

primes = Prime()
print primes[1000]
print primes[2000]
for prime in primes:
    print prime

あまり速くはないが、クラスに機能を加えていけば便利そう。