Python嵌套列表搜索优化

发布于12月02日

我有一个搜索和测试问题: 在从2到100k的素数列表中，我们使用以下条件搜索第一组5:

P1 & lt ; p5 )
解中2个素数的任何组合(3和7=&gt；37和73)也必须是素数
SUM(p1..p5)是满足条件的素数的最小可能和，且高于 100k

我完全可以编写这样的代码，但我有一个严重的优化问题:我的代码超级慢.我有一个低于100k的素数列表，一个超过100k的素数列表，以及一个工作得很好的素数测试，但我不知道如何优化它以在正确的时间获得结果.

对于一个基本的 idea :

100k以下的所有素数列表包含9592项
10亿以下所有素数的列表包含大约5100万行
我有一张长度在10亿以下的素数 list

谢谢你的帮助

推荐答案

Here is a numba version that computes the minimal combination of the prime numbers with restrictions as stated in the question.

运行时的大部分时间都用来预计算素数的所有有效组合.在我的计算机(AMD 5700X)上，这在1分20秒内运行:

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()])

打印:

[    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