【模板】线性筛素数

题意

给出上界 n=10^8 和若干询问，每次输出第 k 小的素数。

思路

Python 不能用普通整数列表保存一亿个筛标记。代码采用奇数分段埃氏筛：

先筛出不超过 sqrt(n) 的基础素数；
每次只处理约一百万个奇数 low,low+2,...,high；
对每个奇基础素数，把块内倍数对应的切片批量置零；
用 bytearray.find(1) 在底层查找仍为素数的位置；
只保存到询问所需的最大排名即可。

块内下标 index 对应整数 low+2*index。若质数 p 的第一个奇倍数位置是 index，后续奇倍数下标相差 p，因此可以写成 flags[index::p]=0。

这种实现没有第三方依赖，完整上限实测约使用 36 MB 内存。

Python 知识

bytearray 每个标记仅一字节，扩展切片赋值在底层批量完成筛除。
array('I') 以 4 字节无符号整数保存约 576 万个素数，显著小于 Python 整数列表。
bytearray.find(1,start) 避免逐字节 Python 循环寻找未标记位置。
只筛奇数，数值跨度与数组下标之间用 number=low+2*index 转换。
/home/rainboy/mycode/hugo-blog/content/program_language/python/cpp_to_python_pitfalls.md：Python 对象内存和底层切片性能。
/home/rainboy/mycode/hugo-blog/content/program_language/python/oj_input_output_cheatsheet.md：大量查询批量输出。

代码

python

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57import sys
from array import array


ODD_BLOCK_SIZE = 1 << 20


def small_primes(limit):
    sieve = bytearray(b"\x01") * (limit + 1)
    sieve[:2] = b"\x00\x00"
    for prime in range(2, int(limit ** 0.5) + 1):
        if sieve[prime]:
            start = prime * prime
            sieve[start::prime] = b"\x00" * ((limit - start) // prime + 1)
    return [number for number in range(3, limit + 1, 2) if sieve[number]]


def first_primes(limit, needed):
    primes = array("I", [2])
    base_primes = small_primes(int(limit ** 0.5))

    low = 3
    while low <= limit and len(primes) < needed:
        high = min(limit, low + 2 * ODD_BLOCK_SIZE - 2)
        if high % 2 == 0:
            high -= 1
        size = (high - low) // 2 + 1
        is_prime = bytearray(b"\x01") * size

        for prime in base_primes:
            if prime * prime > high:
                break
            start = max(prime * prime, (low + prime - 1) // prime * prime)
            if start % 2 == 0:
                start += prime
            index = (start - low) // 2
            is_prime[index::prime] = b"\x00" * ((size - 1 - index) // prime + 1)

        index = is_prime.find(1)
        while index != -1 and len(primes) < needed:
            primes.append(low + 2 * index)
            index = is_prime.find(1, index + 1)
        low = high + 2

    return primes


def main():
    data = list(map(int, sys.stdin.buffer.read().split()))
    limit = data[0]
    queries = data[2:]
    primes = first_primes(limit, max(queries))
    print("\n".join(str(primes[index - 1]) for index in queries))


if __name__ == "__main__":
    main()

复杂度

筛法时间复杂度约 $O(n\log\log n)$ ，查询 $O(1)$ ；空间为基础素数、一个分块和所需质数表，约 $O(\sqrt n+B+\pi(n))$ 。

总结

同一个算法在 Python 中还要重新设计存储。分段、奇数压缩、bytearray 切片和紧凑 array 共同解决了一亿范围的时间与内存问题。

一图流解析

下面保留原题解已有的复盘图片，本轮未重新生成。

【模板】线性筛素数

目录

题意

思路

Python 知识

代码

复杂度

总结

一图流解析