久々にProject Eulerを解いていて素数の判定をしないといけなくなったのでメモ。簡単な問題でも素数の判定が入っていると結構時間がかかるのでエストラネスの篩を使うべき。

普通のプログラム

def isprime(n):
    if n <= 1:
        return False
    elif n == 2:
        return True
    elif n % 2 == 0:
        return False
    for i in range(3, int(n ** 0.5) + 1, 2):
        if n % i == 0:
            return False
    else:
        return True

これの計算時間は$O(\sqrt{n})$

エストラネスの篩を用いる

def isprime(n):
    #
    if n in (2, 3):
        return True
    if n <= 1 or (n % 6 in (0, 2, 3, 4)):
        return False
    #
    import numpy as np
    sieve = np.ones(n // 3, dtype=np.bool)
    for i in range(int(n ** 0.5) // 3 + 1):
        if sieve[i]:
            prime = (3 * i + 4) | 1
            step_ind = 2 * prime
            sieve[prime // 3 + step_ind - 1::step_ind] = False
            sieve[(5 * prime) // 3 - 1::step_ind] = False
    return sieve[-1]

素数定理によると素数の割合は$\dfrac{1}{\log n}$くらい。素数の個数分だけリストを操作するので$O(\dfrac{n}{\log n})$になる。

もう少し改善

この関数を何度も呼ぶのは非効率である。そもそも数回しか呼ばないのなら簡単なプログラムを使ってもいい。そこで素数の大きさの上限が分かっているときと無いときに分けて考える。

必要な素数の大きさの上限が分かっているとき

必要な分だけ素数のリストを作ってクロージャで関数として返す。リストの操作には$O(1)$しかからない。

def isprime_not_over(n):
    if n > 4:
        import numpy as np
        sieve = np.ones(n // 3, dtype=np.bool)
        for i in range(int(n ** 0.5) // 3 + 1):
            if sieve[i]:
                prime = (3 * i + 4) | 1
                step_ind = 2 * prime
                sieve[prime//3 + step_ind - 1::step_ind] = False
                sieve[(5 * prime)//3 - 1::step_ind] = False

    def is_prime(k):
        if k > n:
            raise ValueError("input size over %d" % n)
        elif k in (2, 3):
            return True
        elif k <= 1 or (k % 6 in (0, 2, 3, 4)):
            return False
        else:
            return sieve[k // 3 - 1]

    return isprime

is_prime = isprime_not_over(1000)

必要な素数の大きさの上限が分からないとき

whileとかでループを回していて必要な範囲の上限が分からない場合。

デフォルト引数を使う。

def isprime(n, sieve=[], max=4):
    #
    if n in (2, 3):
        return True
    elif n <= 1 or (n % 6 in (0, 2, 3, 4)):
        return False
    #
    elif n <= max:
        return sieve[n // 3 - 1]
    #
    import numpy as np
    sieve = np.ones((2 * n) // 3, dtype=np.bool)
    for i in range(int(n**0.5) // 3 + 1):
        if sieve[i]:
            prime = (3 * i + 4) | 1
            step_ind = 2 * prime
            sieve[prime//3 + step_ind - 1::step_ind] = False
            sieve[(5 * prime)//3 - 1::step_ind] = False
    max = 2 * n
    return sieve[-1]

計算時間は上限が分かっている時よりも伸びる。maxの変更の仕方を問題に合わせて変えてもいいかもしれない?。もう少し良い方法があるようにも思う。

素数のリストが欲しい時

def primesfrom2under(n):
    sieve = np.ones(n // 3 + (n % 6 == 2) - 1, dtype=np.bool)
    for i in range(int(n ** 0.5) // 3 + 1):
        if sieve[i]:
            prime = (3 * i + 4) | 1
            step_ind = 2 * prime
            sieve[prime//3 + step_ind - 1::step_ind] = False
            sieve[(5 * prime)//3 - 1::step_ind] = False
    primes_over_5 = (3 * np.nonzero(sieve)[0] + 4) | 1
    primes = np.hstack([[2,3], primes_over_5])
    return primes

参考

stackoverflow: fastest-way-to-list-all-primes-below-n-in-pythonを参考にした。エストラネスの篩のところは自分のプログラムのほうが明確だと思う。

tak0kadaの何でもノート

発声練習、生存確認用。

エストラネスの篩を使って素数判定する