オットセイの経営日誌

データサイエンス系ベンチャーを経営してます。経営のこと、趣味のことつぶやきます。

LeetCode / Count Primes

https://leetcode.com/problems/count-primes/

Count the number of prime numbers less than a non-negative number, n.

Example:

Input: 10
Output: 4
Explanation: There are 4 prime numbers less than 10, they are 2, 3, 5, 7.

与えられた値より小さい素数の数を返す問題。

解答・解説

誤答(TLE)

TLEですが、重要な考え方その1が含まれます。それは、

nが何らかの数pで割り切れる場合、n=pqであり、qがpより小さい場合には既にqもしくはqの約数で確かめた際に素因数が検出されているはずである。したがって、素因数候補として確かめるべきは{\displaystyle {\sqrt {n}}}までで十分

です。

以上の考え方を取り入れたのが下記コードで、偶数の判定を飛ばす(2以外素数でない)改善も加えています。

def is_prime(n):
    if n <= 1:
        return False
    if n == 2:
        return True
    if n % 2 == 0:
        return False
    for i in range(3, math.floor(math.sqrt(n))+1, 2):
        if n % i == 0:
            return False
    return True

class Solution:
    def countPrimes(self, n: int) -> int:
        return sum([is_prime(i) for i in range(n)])

が、これだとまだTLEになります。もう一段の高速化を考えないといけません。

解法1

さらに高速化するにあたり、エラトステネスの篩というアルゴリズムがあります。これは、

{\displaystyle x^{1/2}} 以下の素数が既知のとき、{\displaystyle x^{1/2}} 以上 x 以下の素数を決定するには、x 以下の整数で {\displaystyle x^{1/2}} 以下の素数の倍数を全て取り除けば(= 篩えば)よい

というものです。

class Solution:
    def countPrimes(self, n):
        if n < 3:
            return 0
        primes = [True] * n
        primes[0] = primes[1] = False
        for i in range(2, int(n ** 0.5) + 1):
            if primes[i]:
                primes[i * i: n: i] = [False] * len(primes[i * i: n: i])
        return sum(primes)