当前位置：首页 > 文章列表 > 文章 > python教程 > Python嵌套列表素数高效查找方法

Python嵌套列表素数高效查找方法

2025-09-12 08:54:55 0浏览收藏

学习文章要努力，但是不要急！今天的这篇文章《Python嵌套列表素数搜索优化》将会介绍到等等知识点，如果你想深入学习文章，可以关注我！我会持续更新相关文章的，希望对大家都能有所帮助！

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尤其适用于数值计算密集型的代码，例如本例中的素数搜索。

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