本文旨在提供一种优化Python代码，以解决在素数列表中搜索满足特定条件的素数组合的问题。通过使用Numba库进行即时编译，并结合并行计算，可以显著提高搜索效率。本文将详细介绍如何使用Numba优化代码，并提供完整的示例代码。
问题描述
我们需要在一个包含2到10万的素数列表中，找到满足以下条件的第一个包含5个素数的集合：
- p1
- 该集合中任意两个素数组合（例如，3和7 => 37和73）也必须是素数
- sum(p1..p5) 是满足上述条件的素数和的最小值，且大于10万
原始代码在解决这个问题时速度非常慢，因此需要进行优化。
优化方案：使用Numba
Numba是一个Python的即时（JIT）编译器，它可以将Python代码转换为机器码，从而显著提高代码的执行速度。Numba尤其适用于数值计算密集型的代码，例如本例中的素数搜索。
立即学习“Python免费学习笔记（深入）”；
1. 安装Numba
首先，确保已经安装了Numba。可以使用pip进行安装：
pip install numba
2. 代码优化步骤
以下是使用Numba优化代码的步骤：
a. 导入必要的库
import numpy as np from numba import njit, prange
- numpy 用于数组操作。
- numba.njit 用于将Python函数编译为机器码。
- numba.prange 用于并行循环。
b. 定义素数判断函数
@njit
def prime(a):
if a < 2:
return False
for x in range(2, int(a**0.5) + 1):
if a % x == 0:
return False
return True使用@njit装饰器告诉Numba编译此函数。
c. 定义字符串转整数函数
@njit
def str_to_int(s):
final_index, result = len(s) - 1, 0
for i, v in enumerate(s):
result += (ord(v) - 48) * (10 ** (final_index - i))
return result此函数用于将两个素数连接成一个整数，例如将3和7连接成37。
d. 生成素数列表
@njit
def generate_primes(n):
out = []
for i in range(3, n + 1):
if prime(i):
out.append(i)
return out生成小于n的所有素数。
e. 生成有效的素数组合
@njit(parallel=True)
def get_comb(n=100_000):
# generate all primes < n
primes = generate_primes(n)
n_primes = len(primes)
# generate all valid combinations of primes
combs = np.zeros((n_primes, n_primes), dtype=np.uint8)
for i in prange(n_primes):
for j in prange(i + 1, n_primes):
p1, p2 = primes[i], primes[j]
c1 = str_to_int(f"{p1}{p2}")
c2 = str_to_int(f"{p2}{p1}")
if not prime(c1) or not prime(c2):
continue
combs[i, j] = 1
all_combs = []
for i_p1 in prange(0, n_primes):
for i_p2 in prange(i_p1 + 1, n_primes):
if combs[i_p1, i_p2] == 0:
continue
for i_p3 in prange(i_p2 + 1, n_primes):
if combs[i_p1, i_p3] == 0:
continue
if combs[i_p2, i_p3] == 0:
continue
for i_p4 in prange(i_p3 + 1, n_primes):
if combs[i_p1, i_p4] == 0:
continue
if combs[i_p2, i_p4] == 0:
continue
if combs[i_p3, i_p4] == 0:
continue
for i_p5 in prange(i_p4 + 1, n_primes):
if combs[i_p1, i_p5] == 0:
continue
if combs[i_p2, i_p5] == 0:
continue
if combs[i_p3, i_p5] == 0:
continue
if combs[i_p4, i_p5] == 0:
continue
p1, p2, p3, p4, p5 = (
primes[i_p1],
primes[i_p2],
primes[i_p3],
primes[i_p4],
primes[i_p5],
)
ccomb = np.array([p1, p2, p3, p4, p5], dtype=np.int64)
if np.sum(ccomb) < n:
continue
all_combs.append(ccomb)
print(ccomb)
break
return all_combs- @njit(parallel=True) 启用并行计算。
- prange 用于并行循环，可以显著提高计算速度。
- 首先，生成所有素数小于n的列表。
- 然后，生成所有有效的素数组合，存储在combs数组中。
- 最后，遍历所有可能的素数组合，找到满足条件的第一个组合。
f. 找到最小的素数组合
all_combs = np.array(get_comb())
print()
print("Minimal combination:")
print(all_combs[np.sum(all_combs, axis=1).argmin()])计算所有素数组合的和，并找到最小的和对应的组合。
3. 完整代码
import numpy as np
from numba import njit, prange
@njit
def prime(a):
if a < 2:
return False
for x in range(2, int(a**0.5) + 1):
if a % x == 0:
return False
return True
@njit
def str_to_int(s):
final_index, result = len(s) - 1, 0
for i, v in enumerate(s):
result += (ord(v) - 48) * (10 ** (final_index - i))
return result
@njit
def generate_primes(n):
out = []
for i in range(3, n + 1):
if prime(i):
out.append(i)
return out
@njit(parallel=True)
def get_comb(n=100_000):
# generate all primes < n
primes = generate_primes(n)
n_primes = len(primes)
# generate all valid combinations of primes
combs = np.zeros((n_primes, n_primes), dtype=np.uint8)
for i in prange(n_primes):
for j in prange(i + 1, n_primes):
p1, p2 = primes[i], primes[j]
c1 = str_to_int(f"{p1}{p2}")
c2 = str_to_int(f"{p2}{p1}")
if not prime(c1) or not prime(c2):
continue
combs[i, j] = 1
all_combs = []
for i_p1 in prange(0, n_primes):
for i_p2 in prange(i_p1 + 1, n_primes):
if combs[i_p1, i_p2] == 0:
continue
for i_p3 in prange(i_p2 + 1, n_primes):
if combs[i_p1, i_p3] == 0:
continue
if combs[i_p2, i_p3] == 0:
continue
for i_p4 in prange(i_p3 + 1, n_primes):
if combs[i_p1, i_p4] == 0:
continue
if combs[i_p2, i_p4] == 0:
continue
if combs[i_p3, i_p4] == 0:
continue
for i_p5 in prange(i_p4 + 1, n_primes):
if combs[i_p1, i_p5] == 0:
continue
if combs[i_p2, i_p5] == 0:
continue
if combs[i_p3, i_p5] == 0:
continue
if combs[i_p4, i_p5] == 0:
continue
p1, p2, p3, p4, p5 = (
primes[i_p1],
primes[i_p2],
primes[i_p3],
primes[i_p4],
primes[i_p5],
)
ccomb = np.array([p1, p2, p3, p4, p5], dtype=np.int64)
if np.sum(ccomb) < n:
continue
all_combs.append(ccomb)
print(ccomb)
break
return all_combs
all_combs = np.array(get_comb())
print()
print("Minimal combination:")
print(all_combs[np.sum(all_combs, axis=1).argmin()])4. 运行结果
在AMD 5700X CPU上，该代码可以在1分20秒内完成计算。
[ 3 28277 44111 70241 78509] [ 7 61 25939 26893 63601] [ 7 61 25939 61417 63601] [ 7 61 25939 61471 86959] [ 7 2467 24847 55213 92593] [ 7 3361 30757 49069 57331] ... [ 1993 12823 35911 69691 87697] [ 2287 4483 6793 27823 67723] [ 3541 9187 38167 44257 65677] Minimal combination: [ 13 829 9091 17929 72739] real 1m20,599s user 0m0,011s sys 0m0,008s
注意事项
- Numba在第一次运行时需要编译代码，因此第一次运行可能会比较慢。
- Numba对某些Python特性支持有限，需要注意代码的兼容性。
- 并行计算可以提高计算速度，但也会增加CPU的负载。
总结
通过使用Numba进行即时编译和并行计算，可以显著提高Python代码的执行速度，特别是在数值计算密集型的任务中。本文提供了一个具体的例子，展示了如何使用Numba优化素数搜索问题，并提供了完整的示例代码。希望本文能帮助读者更好地理解和应用Numba。
