低内存场景下双分段素数筛函数的性能提速方法问询

API 探路者

2026-8-28

在可用内存受限的场景下，可通过实现双分段素数筛（Double Segmented Primes Sieve）来查找n邻域内的素数。实际实现中，分段筛设置维度为dim_seg_1的筛区处理小于sqrt(n)的小素数，设置大小为dim_seg_2的筛区处理更大数值，以下为用于查找区间[n_start, n_start + dim_seg_2]内素数的函数实现：

import numpy as np

def double_segmented_sieve(dim_seg_1,dim_seg_2,n_start):
    dim1=dim_seg_1//6
    dim2=dim_seg_2//6+1
    kmin=n_start//6-1
    sqrt_nmax=int(((n_start+dim_seg_2)**0.5+1)/6)+1
    if sqrt_nmax>(dim_seg_1**2)//6:
        return -1

    # 筛区间[n_start, n_start+dim_seg_2]内素数所用数组
    sieve5m6 =np.ones((dim2+1), dtype=bool)
    sieve1m6 =np.ones((dim2+1), dtype=bool)
    
    # 筛sqrt(n)范围内小素数所用数组
    sieve5m6_1 =np.ones((dim1+1), dtype=bool)
    sieve1m6_1 =np.ones((dim1+1), dtype=bool)
    for i in range(1,int((dim_seg_1**0.5+1)/6)+1):
        if sieve5m6_1[i]:
            sieve5m6_1[6*i*i::6*i-1]=[False]
            sieve1m6_1[6*i*i-2*i::6*i-1]=[False]
        if sieve1m6_1[i]:
            sieve5m6_1[6*i*i::6*i+1]=[False]
            sieve1m6_1[6*i*i+2*i::6*i+1]=[False]
    sieve5m6_2 =sieve5m6_1
    sieve1m6_2 =sieve1m6_1 
  
    for j in range(0,sqrt_nmax//dim1):
        for i in range(1,dim1+1):
            if sieve5m6_2[i]:
                sieve5m6[(-kmin+i+j*dim1)%(6*(i+j*dim1)-1)::6*(i+j*dim1)-1]=[False]
                sieve1m6[(-kmin-i-j*dim1)%(6*(i+j*dim1)-1)::6*(i+j*dim1)-1]=[False]
            if sieve1m6_2[i]:
                sieve5m6[(-kmin-i-j*dim1)%(6*(i+j*dim1)+1)::6*(i+j*dim1)+1]=[False]
                sieve1m6[(-kmin+i+j*dim1)%(6*(i+j*dim1)+1)::6*(i+j*dim1)+1]=[False]
        
        sieve5m6_2 =np.ones((dim1+1), dtype=bool)
        sieve1m6_2 =np.ones((dim1+1), dtype=bool)
        for i in range(1,int(((-1+(1+6*(j+1)*dim1))**0.5)/6)+1):
            if sieve5m6_1[i]:
                sieve5m6_2[(-(j+1)*dim1+i)%(6*i-1)::6*i-1]=[False]
                sieve1m6_2[(-(j+1)*dim1-i)%(6*i-1)::6*i-1]=[False]
            if sieve1m6_1[i]:
                sieve5m6_2[(-(j+1)*dim1-i)%(6*i+1)::6*i+1]=[False]
                sieve1m6_2[(-(j+1)*dim1+i)%(6*i+1)::6*i+1]=[False]

    sieve5m6[0:1:]=[False]  
    sieve1m6[0:0:]=[False]
    return np.r_[ (6 *(kmin+ np.nonzero(sieve5m6)[0]) - 1),(6 *(kmin+  np.nonzero(sieve1m6)[0]) + 1)]

P=np.sort(double_segmented_sieve(2**26,2**14,2**60))

请问可以通过哪些方式对该函数的运行速度进行优化？

优化方案

现有实现已经通过模6简化将筛内存占用压缩到了传统埃氏筛的1/3，但仍有较大提速空间，可按优先级从以下方向调整：

提取小素数列表，减少无效遍历：当前分段筛大区间时，每轮都要遍历长度为dim_seg_1//6的布尔数组判断位置是否对应素数，其中大部分位置都是非素数，无效判断占比很高。可以在第一次完成小素数区间筛选后，直接用np.nonzero把所有小素数的数值提取成单独的一维数组，后续分段筛时直接遍历这个素数列表即可，能砍掉90%以上的内层循环判断开销。
修正Numpy操作的低效写法：当前所有切片赋值都用=[False]的写法，会先生成单元素Python列表再做广播，比直接赋值标量False慢30%左右，直接改成=False即可。另外同一分支内重复计算的6*(i+j*dim1)-1、6*(i+j*dim1)+1这类素数值、偏移量，可以提前存为临时变量，避免重复做算术运算。
替换取模运算为递推偏移：当前每处理一个素数都用取模计算它在当前分段的第一个标记位置，取模属于开销较高的运算。可以在初始化时给每个小素数算好在第一个大分段的起始偏移，之后每处理完一个分段，直接把偏移量减去当前分段长度，小于0就加上素数本身的步长，完全不需要取模，运算开销能降一半以上。
使用位级筛数组提升缓存命中率：当前用的np.bool_类型每个元素占1字节，内存利用率很低。可以换成位打包数组，比如用np.uint8类型每个元素存8个标记位，筛数组内存直接降到原来的1/8，CPU缓存命中率会大幅提升，大尺寸分段下速度能提升2~3倍。
减少不必要的内存申请和拷贝：当前每轮j循环都会重新调用np.ones初始化两个小筛数组，反复申请释放内存的开销很高。可以在循环外提前初始化这两个小数组，每轮处理完只需要把上一轮标记为False的位置重置为True即可，不需要全量重新赋值。另外最后返回结果时用的np.r_是动态拼接，效率很低，可以提前计算两个素数集合的总长度，预分配好结果数组再填充，比动态拼接快不少。
把Python层循环下放到C层：当前最内层的素数标记循环是跑在Python解释器上的，哪怕切片用了Numpy，循环本身的开销仍然占大头。可以把所有小素数的步长、起始偏移整理成Numpy数组，批量生成切片索引做标记，把循环逻辑交给Numpy的C实现处理，速度能提升一个量级。如果追求更高性能，给核心循环加上Numba的@njit装饰器做JIT编译，基本不用改逻辑就能拿到接近原生C的速度。
清理无效代码：当前代码里sieve1m6[0:0:] = [False]是空切片操作，不会产生任何实际效果，可以直接删掉。开头的平方根阈值、循环边界值可以提前一次性算好，不要在循环里重复计算。

内容的提问来源于stack exchange，提问作者user140242