发布于 

使用 Python 实现素数筛时发现的有趣问题

最近学习素数筛算法，算法参考自两篇 CSDN 的文章12，原文是用 C/C++ 来写的，我用 Python 进行了改写。

1
2
3
4
5
import time
from icecream import ic
import numba
from numba import jit
import numpy as np

这里我使用了 icecream 模块进行调试。

先简单介绍一下素数筛算法。对于一个给定的自然数序列，要找出其中的所有素数，可以通过筛除合数的方法来寻找。筛法中思路最简单的是埃氏（Eratosthenes）筛，基本思路是：若要寻找 n 以内的所有素数，只需找不大于 n1/2n^{1/2} 的素数，剔除它们的 2 倍或 2 倍以上的数，从小到大，依次剔除，即可不用单独判断素数，得到素数表。但是这种算法存在一个缺陷，即很多数的质因数不止一个，所以它们被重复判断了。于是有了改进后的线性筛，即欧拉（Euler）筛。欧拉筛的原理是对于每次筛除，只用它的最小质因数筛去，即把这个数表示为一个较大数的质数倍数，这样每个数只用筛去一次，消除了重复。

以下是我从 C 改写的 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
def Eratosthenes(n):
    primes = [1 for _ in range(n+1)]
    primes[0] = primes[1] = 0
    count = 0
    for i in range(2, n + 1):
        if primes[i]:
            for j in range(i * 2, n + 1, i):
                primes[j] = 0
                count += 1
    c = 0
    for i in primes:
        if i:
            c += 1
    return count,c
  
def Euler(n):
    primes = [1 for _ in range(n+1)]
    primes[0] = primes[1] = 0
    res = []
    count = 0
    for i in range(2, n + 1):
        if primes[i]:
            res.append(i)     
        j = 0
        index = res[j] * i
        while j < len(res) and index <= n:
            count += 1
            primes[index] = 0
            if i % res[j] == 0:
                break
            else:
                j += 1
                index = res[j] * i
    return count,len(res)

其中第一个返回值是筛除次数，第二个返回值是素数个数。

由于 Python 中数组的处理可以使用 Numpy，所以我又对代码进行了改进，得到了一系列函数：

首先，对于 Eratosthenes 方法，从 2i2i 开始显然是低效的，改为从 i2i^2 开始：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def Eratosthenes_improved(n):
    primes = [1 for _ in range(n+1)]
    primes[0] = primes[1] = 0
    count = 0
    for i in range(2, n + 1):
        if primes[i]:
            for j in range(i * i, n + 1, i):
                primes[j] = 0
                count += 1
    c = 0
    for i in primes:
        if i:
            c += 1
    return count,c

考虑使用 Numpy。事实上，之后的实验发现，不恰当的使用 Numpy 并不会提升速度。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
def Eratosthenes_numpy(n):
    primes = np.ones(n + 1)
    primes[0] = primes[1] = 0
    count = 0
    for i in range(2, int(n ** 0.5) + 1):
        if primes[i]:
            for j in range(i * i, n + 1, i):
                primes[j] = 0
                count += 1
    c = 0
    for i in primes:
        if i:
            c += 1
    return count,c

对于 Euler 方法，append 的使用会影响速度。事实上，在埃氏筛的算法中，缺少了生成素数列表的操作，额外的实验表明 ，这些操作所占的时间并不会影响整体的结果。不过，这里还是对 append 进行了改进，另外也消除了 len 对于速度的影响：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def Euler_improved(n):
    primes = [1 for _ in range(n + 1)]
    primes[0] = primes[1] = 0
    res = [1 for _ in range(int(n / 2))]
    count = 0
    c = 0
    for i in range(2, n + 1):
        if primes[i]:
            res[c] = i
            c += 1  
        j = 0
        index = res[j] * i
        while j < c and index <= n:
            count += 1
            primes[index] = 0
            if i % res[j] == 0:
                break
            else:
                j += 1
                index = res[j] * i
    return count

考虑使用 Numpy 的情况，首先是保留append，之后的实验表明，Numba 的加速会产生相反的结果：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def Euler_numpy(n):
    primes = np.ones(n + 1)
    primes[0] = primes[1] = 0
    res = []
    count = 0

    for i in range(2, n + 1):
        if primes[i]:
            res.append(i)
        j = 0
        index = res[j] * i
        while j < len(res) and index <= n:
            count += 1
            primes[index] = 0
            if np.mod(i, res[j]) == 0:
                break
            else:
                j += 1
                index = res[j] * i
    return count

接下来是改进的情况：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def Euler_improved_numpy(n):
    primes = np.ones(n + 1)
    primes[0] = primes[1] = 0
    res = np.zeros(int(n / 2))
    count = 0
    c = 0
    for i in range(2, n + 1):
        if primes[i]:
            res[c] = i
            c += 1
        
        j = 0
        index = int(res[j] * i)
        while j < c and index <= n:
            count += 1
            primes[index] = 0
            if i%res[j] == 0:
                break
            else:
                j += 1
                index = int(res[j] * i)
    return count

观察发现，Euler 方法中，每次循环多出大数取模运算和乘法运算。于是简单测量了一下时长：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
s = []
for _ in range(10):
    lst = [0 for _ in range(10000000)]
    t0 = time.perf_counter()
    lst[5000000] = 1
    t1 = time.perf_counter()
    dt1 = t1 - t0
    lst = [0 for _ in range(10000000)]
    t0 = time.perf_counter()
    lst[5000000] = 1
    a = 399999 % 19
    b = 19 * 399999
    t1 = time.perf_counter()
    dt2 = t1 - t0
    ic(f'{dt2/dt1:.2f}')
    s.append(dt2/dt1)
ic(sum(s)/len(s))

执行多次以后发现，时长比有着一定范围的波动，最高可到 2 倍以上：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
ic| f'{dt2/dt1:.2f}': '2.20'
ic| f'{dt2/dt1:.2f}': '1.37'
ic| f'{dt2/dt1:.2f}': '1.44'
ic| f'{dt2/dt1:.2f}': '1.09'
ic| f'{dt2/dt1:.2f}': '1.32'
ic| f'{dt2/dt1:.2f}': '1.16'
ic| f'{dt2/dt1:.2f}': '1.33'
ic| f'{dt2/dt1:.2f}': '1.22'
ic| f'{dt2/dt1:.2f}': '1.33'
ic| f'{dt2/dt1:.2f}': '1.18'
ic| sum(s)/len(s): 1.3642452780259806

ic| f'{dt2/dt1:.2f}': '1.46'
ic| f'{dt2/dt1:.2f}': '1.20'
ic| f'{dt2/dt1:.2f}': '1.18'
ic| f'{dt2/dt1:.2f}': '1.50'
ic| f'{dt2/dt1:.2f}': '1.00'
ic| f'{dt2/dt1:.2f}': '1.69'
ic| f'{dt2/dt1:.2f}': '1.33'
ic| f'{dt2/dt1:.2f}': '1.33'
ic| f'{dt2/dt1:.2f}': '1.62'
ic| f'{dt2/dt1:.2f}': '1.18'
ic| sum(s)/len(s): 1.3493646302728954

平均水平在 1.3 左右。

所以采用了 Numpy 中的取模运算：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def Euler_improved_numpy_with_mod(n):
    primes = np.ones(n + 1)
    primes[0] = primes[1] = 0
    res = np.zeros(int(n / 2))
    count = 0
    c = 0
    for i in range(2, n + 1):
        if primes[i]:
            res[c] = i
            c += 1
        
        j = 0
        index = int(res[j] * i)
        while j < c and index <= n:
            count += 1
            primes[index] = 0
            if np.mod(i,res[j]) == 0:
                break
            else:
                j += 1
                index = int(res[j] * i)
    return count

我从 StackOverflow 上搜索有关问题，又对埃氏筛进行了改进，同时发现有大佬给出了几种超快的算法3，我选取了其中两种，一种是纯 Python 的，另一种是使用 Numpy 的：

1
2
3
4
5
6
7
8
9
10
11
def Eratosthenes_numpy_improved(n):
    primes = np.ones(n + 1)
    primes[0] = primes[1] = 0
    for i in range(2, int(n ** 0.5) + 1):
        if primes[i]:
            primes[i * i::i] = 0
    c = 0
    for i in primes:
        if i:
            c += 1
    return c

使用 Numpy 的切片运算可以显著提高速度，Stackoverflow 上的方法也用到了切片运算：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
def primes(n):
    sieve = [True] * n
    for i in range(3,int(n**0.5)+1,2):
        if sieve[i]:
            sieve[i*i::2*i]=[False]*((n-i*i-1)//(2*i)+1)
    res =  [2] + [i for i in range(3,n,2) if sieve[i]]
    return len(res)
    
def primesfrom2to(n):
    sieve = np.ones(n//3 + (n%6==2), 
                    dtype=bool)
    for i in range(1,int(n**0.5)//3+1):
        if sieve[i]:
            k=3*i+1|1
            sieve[k*k//3::2*k] = False
            sieve[k*(k-2*(i&1)+4)//3::2*k] = False
    res =  np.r_[2,3,((3*np.nonzero(sieve)[0][1:]+1)|1)]
    return len(res)

使用不同的数量级对这些函数的速度进行测试，发现大多数函数运行时间基本是 10 倍变化，这是合理的:

1
2
3
4
5
6
7
# 测试函数
def prime_list(func, n):
    t0 = time.perf_counter()
    res = func(n)
    t1 = time.perf_counter()
    dt = t1 - t0
    return dt, res

时间结果附在文后4。

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
# 绘图代码
import pandas as pd
import matplotlib.pyplot as plt
from icecream import ic
import numpy as np

def to_float(series):
    series_new = pd.Series(dtype='float64')
    for i in range(len(series)):
        series_new[str(i)] = np.log(float(series.iloc[i]))+14
    argmax = int(series_new.idxmax())
    argmin = int(series_new.idxmin())
    return {'y':series_new,'idxmax':argmax,'max':series_new[argmax],
            'idxmin':argmin,'min':series_new[argmin]}

df = pd.read_csv(' 素数筛速度统计.csv',dtype = 'object').\\
                  drop(columns=['numba(jit no python)'],axis=1)
nojit = df[0:11]
funcs = df['function'][0:11]
plt.figure(figsize=(10,10))
for i in range(8,0,-1):
    res = to_float(nojit.loc[:,str(10 ** i)])
    plt.barh(funcs,res['y'])
plt.yticks(fontsize = 10)
plt.subplots_adjust(left=0.25)
plt.legend(loc = 'upper right',labels = [10 ** i for i in range(8,0,-1)])
plt.savefig('nojit.svg')
plt.show()

得到图像：

这里对运行时间进行了处理，来更方便地展示：

t′=ln⁡t+14t'=\ln t+14

可以看到，在参考 StackOverflow 前的方法中，仅仅是使用 Numpy 生成数组并不能加快速度，反而会变得很慢。同时，虽然 Euler 方法减少了重复，但实际运行速度并没有加快，可能有循环步骤中操作过多的原因。而且在小规模时，使用 Numpy 的 Eratosthenes numpy improved 和 primesfrom2to 并没有表现出优势，大规模时优势才凸显出来，而使用原生列表的 primes 始终保持着较好的速度，这可能是 Numpy 的原因。

在使用 Numba 进行加速以后，得到新的运行时间：

1
2
# 在每个函数前添加
 @jit(nopython = True)

需要说明的是，由于 primesfrom2to 无法完全脱离 Python 解释器运行，所以使用了：

1
@jit(）

使用同样的方法进行测试并绘制图像：

1
2
3
4
5
6
7
8
9
10
11
12
13
jit = df[11:]
plt.figure(figsize=(10,10))
for i in range(8,0,-1):
    res = to_float(jit.loc[:,str(10 ** i)])
    if i == 1:
        plt.barh(funcs,res['y'],alpha = 0.3)
    else:
        plt.barh(funcs,res['y'])
plt.yticks(fontsize = 10)
plt.subplots_adjust(left=0.25)
plt.legend(loc = 'upper right',labels = [10 ** i for i in range(8,0,-1)])
plt.savefig('jit.svg')
plt.show()

这是由于在规模为 10 时，速度较慢，呈现在图中会严重覆盖其他内容，所以调整了它的透明度，图像中变暗的地方时规模为 10 的情况：

可以看到，Numba 的加速效果是显而易见的，速度最慢是 16 左右。由于 primesfrom2to 基本没有变化，而且本身速度很快，所以这里不考虑。考虑其他函数，发现加速后 Numpy 的优势表现出来，而且经过加速，Euler 方法在大规模问题上优于 Eratosthenes 方法。有趣的是，加速以后 Euler 的 improved 方法表现变差，可能是 Numba 对于 append 操作有着较好的优化的原因，比我自己简单的改进更高效。

总的来看，使用 Eratosthenes numpy improved 或者 Numba 加速后的 Euler numpy 就可以在 1 亿规模获得简单好用的效果。

许可协议

本文采用署名-非商业性使用-相同方式共享 4.0 国际许可协议，转载请注明出处。

分享文章