高速で正しく動くエラトステネスの篩をPythonで書く

はじめに

「エラトステネスの篩 python」で検索すると色々実装されたものが出てきますが、一番上に出てくるページのコードは遅い実装である上にコード自体が間違っているという始末です(どのページかは伏せますがmax=4にすると4が素数判定されるのが1つの反例です)。
自分で早く動く(はずの)ものを実装してみたのでメモ代わりに残しておきます。

追記

更に高速化したものを作成したので記事にしました。

strangerxxx.hateblo.jp

本記事は少し遅いですが競プロで利用するには問題なく、より実装が簡素なのでこのまま置いておきます。

本題

詳しい話はしませんが、appendなどの処理はメモリなどの関係で遅くなるのでなるべく使わないのが吉です。
以下が実装したコードです。

def eratosthenes(limit: int, minLimit: int = None):
    if minLimit and (minLimit < 0 or minLimit > limit):
        raise ValueError("incorrect minLimit")
    isPrime = [True] * max(limit + 1, 2)
    isPrime[0] = False
    isPrime[1] = False
    for p in range(2, limit + 1):
        if not isPrime[p]:
            continue
        for i in range(p * p, limit + 1, p):
            isPrime[i] = False
    if minLimit:
        return [i for i, x in enumerate(isPrime[minLimit:], start=minLimit) if x]
    else:
        return [i for i, x in enumerate(isPrime) if x]

この関数の引数に最大値を指定すれば素数のリストが返ってきます。また第2引数には最小値が指定可能です(省略可)。

print(eratosthenes(19))
# -> [2, 3, 5, 7, 11, 13, 17, 19]
print(eratosthenes(200, 150))
# -> [151, 157, 163, 167, 173, 179, 181, 191, 193, 197, 199]

エラトステネスの篩の計算量は らしいですが、AtCoderのコードテストでは limit=10⁷でもPyPy3で700msほどで実行できました(Python3.8.2では2250msほどでしたが……)。

ジェネレーター版

ついでにジェネレーターにしたものも置いておきます。

def eratosthenes_gen(limit: int, minLimit: int = 0):
    if not 0 <= minLimit <= limit:
        raise ValueError("incorrect minLimit")
    isPrime = [True] * (limit + 1)
    for p in range(2, limit + 1):
        if not isPrime[p]:
            continue
        if p >= minLimit:
            yield p
        for i in range(p * p, limit + 1, p):
            isPrime[i] = False

皆さんもよき素数ライフを！