本文旨在提供一种优化Python代码，以解决在素数列表中搜索满足特定条件的素数组合的问题。通过使用Numba库进行即时编译，并结合并行计算，可以显著提高搜索效率。本文将详细介绍如何使用Numba优化代码，并提供完整的示例代码。
问题描述
我们需要在一个包含2到10万的素数列表中，找到满足以下条件的第一个包含5个素数的集合：
- p1 < p2 < p3 < p4 < p5
- 该集合中任意两个素数组合（例如，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。