如何优化Eratosthenes的NumPy筛?

我已经在NumPy中实现了我自己的埃拉托斯梯尼筛。我相信你们都知道这是为了找到一个数下的所有素数,所以我不做任何进一步的解释。

代码:

import numpy as np

def primes_sieve(n):
    primes = np.ones(n+1, dtype=bool)
    primes[:2] = False
    primes[4::2] = False
    for i in range(3, int(n**0.5)+1, 2):
        if primes[i]:
            primes[i*i::i] = False

    return np.where(primes)[0]

正如你所看到的,我已经做了一些优化,首先,除了2之外,所有素数都是奇数,所以我将2的所有倍数设置为 False 并且只有蛮力奇数。

其次,我只循环遍历数字,直到平方根的底部,因为平方根之后的所有复数都会被平方根以下素数的倍数所消除。

但它不是最优的,因为它循环通过低于极限的所有奇数,并且不是所有奇数都是素数。随着数量的增加,素数变得越来越稀疏,所以有很多冗余的迭代。

因此,如果候选列表是动态变化的,以这种方式,已经识别的复合数甚至永远不会被迭代,因此只有素数循环通过,就不会有任何浪费的迭代,因此算法将是最优的。

我写了一个优化版本的粗略实现:

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]

但它比未优化的版本慢得多:

In [92]: %timeit primes_sieve(65536)
271 µs ± 22 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [102]: %timeit primes_sieve_opt(65536)
309 µs ± 3.86 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

我的想法很简单,通过获取的下一个索引 True ,我可以确保所有素数都被覆盖,并且只处理素数。

然而 np.argmax 在这方面进展缓慢。我在谷歌上搜索了“如何在NumPy数组中找到下一个True值的索引”(没有引号),我发现了几个StackOverflow问题,这些问题稍微相关,但最终没有回答我的问题。

例如 numpy get index where value is true 和 Numpy first occurrence of value greater than existing value 。

我并没有试图找到所有索引 真的 ,这样做是非常愚蠢的,我需要找到下一个 真的 值,获取其索引并立即停止循环,只有 bool s

我如何优化它?

编辑

如果有人感兴趣,我已经进一步优化了我的算法:

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)

我用过 numba 加快速度。由于除了2和3之外的所有素数都是6k+1或6k-1,这使得事情变得更快。