204.计数质数

统计所有小于非负整数 n 的质数的数量。

示例:

输入: 10
输出: 4
解释: 小于 10 的质数一共有 4 个, 它们是 2, 3, 5, 7 。

这个题目思路很简单，但是可能效率不高，里面有一些小trick需要注意，所以在这里记录一下优化过程。

思路1: 逐个判断每个数是否质数, (超时)

class Solution:
    def countPrimes(self, n: int) -> int:
        def isPrime(val):
            if val <= 1:
                return False
            for i in range(2, int(math.sqrt(val)) + 1):
                if val % i == 0:
                    return False
            return True
        count = 0
        for i in range(n):
            if isPrime(i):
                count += 1
        return count

思路2：厄拉多塞筛法

• 2是最小的质数，从2开始遍历
• 如果后面的数是2的倍数，则将其进行标记
• 从下一个未被标记的数开始, 这里是3
• 如果后面的数是3的倍数，则将其进行标记
• 重复上述步骤

实现上的小细节: 可以想到的是i从2-n进行遍历，但怎么判断i是否被标记呢?

1. 被标记的放到list或者dict，每次遇到i进行查询
2. 或者 建立一个size=n数组进行标记

查询操作太慢, 实现会超时，这里采用思路2,

1760 ms 25.3 MB, 超过7.20%

class Solution:
    def countPrimes(self, n: int) ->int:
        nums = [0] * n
        res = 0
        for i in range(2, n):
            if nums[i] == 0:
                mul = 2
                res += 1
                while mul * i < n:
                    nums[mul * i] += 1
                    mul += 1
        return res

再进行一些小优化, while循环换成数组操作

500 ms, 37.2 MB, 超过46.82%

class Solution:
    def countPrimes(self, n: int) ->int:
        nums = [0] * n
        res = 0
        for i in range(2, n):
            if nums[i] == 0:
                res += 1
                nums[2*i : n : i] = [1]*len(nums[2*i : n : i])
        return res

继续优化, 主要是有一个这样的观察:
假如当遍历到5的时候，我们会对 5 * 2, 5 * 3, 5 * 4, 5 * 5, … 进行标记，但其实发现 5 * 2, 5 * 3, 5 * 4, 这些在之前已经被标记了，所以直接从 5 * 5进行标记， 即对于i * i开始遍历; 另外, 当i超过sqrt(n)是， i * i就会超过 n，这时就不会改变标记。 所以只需要遍历到 sqrt(n)。质数的个数就是 n - 被标记的个数。

204 ms, 37.1 MB, 超过77.81%

class Solution:
    def countPrimes(self, n: int) ->int:
        if n < 2:
            return 0
        nums = [0] * n
        nums[1] = nums[0] = 1
        for i in range(2, int(math.sqrt(n)) + 1):
            if nums[i] == 0:
                nums[i*i : n : i] = [1]*len(nums[i*i : n : i])
        return n - sum(nums)