如何优化基于NumPy的埃拉托斯特尼筛法实现？

初步优化的NumPy筛法 #

我用NumPy实现了埃拉托斯特尼筛法，用于查找指定数值以下的所有质数，并完成了两项初步优化：

• 除2外所有质数均为奇数，先将所有2的倍数设为False，仅遍历奇数
• 仅循环至目标数的平方根，因为平方根后的合数会被其平方根以下的质数倍数筛除

尝试仅遍历质数的优化（性能倒退） #

上述实现仍有冗余：会遍历所有奇数（包括合数），数值越大质数越稀疏，冗余迭代越多。我编写了仅遍历质数的优化版本，但因使用np.argmax查找下一个质数，反而导致性能更差：

def primes_sieve_opt(n):
    primes = np.ones(n+1, dtype=bool)
    primes[:2] = False
    primes[4::2] = False
    limit = int(n**0.5)+1
    i = 2
    while i < limit:
        primes[i*i::i] = False
        i += 1 + primes[i+1:].argmax()

    return np.where(primes)[0]

性能测试显示该版本比未优化版慢，我需要一种高效的方法来查找NumPy布尔数组中下一个True值的索引，以此优化筛法的遍历逻辑。

结合Numba与6k±1特性的最终优化 #

后续我进一步优化了算法：用Numba进行加速，同时利用“除2和3外，质数均为6k+1或6k-1形式”的特性，实现如下：

import numba
import numpy as np

@numba.jit(nopython=True, parallel=True, fastmath=True, forceobj=False)
def prime_sieve(n: int) -> np.ndarray:
    primes = np.full(n + 1, True)
    primes[:2] = False
    primes[4::2] = False
    primes[9::6] = False
    limit = int(n**0.5) + 1
    for i in range(5, limit, 6):
        if primes[i]:
            primes[i * i :: 2 * i] = False

    for i in range(7, limit, 6):
        if primes[i]:
            primes[i * i :: 2 * i] = False

    return np.flatnonzero(primes)

内容的提问来源于stack exchange，提问作者Ξένη Γήινος