Python嵌套列表搜索优化：寻找满足特定条件的素数组合-Python教程-PHP中文网

Python嵌套列表搜索优化：寻找满足特定条件的素数组合

霞舞

发布： 2025-09-09 18:37:06

原创

287人浏览过

python嵌套列表搜索优化：寻找满足特定条件的素数组合

本文旨在提供一种优化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优化代码的步骤：

纳米搜索

纳米搜索：360推出的新一代AI搜索引擎

查看详情

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

登录后复制