如何优化查找7个属于阿姆斯特朗数的素数的Python程序

你原来的代码效率低主要是两个核心问题：一是逐一枚举所有整数，大数范围下枚举量爆炸；二是试除法素性判断、重复计算数字幂次的开销太高。以下是可直接落地的优化方案：

核心优化思路 #

• 先限定自幂数范围：n位自幂数的最大可能值为 n * 9^n，当n>60时n * 9^n < 10^(n-1)，不存在更大的自幂数，不用遍历无穷大的范围
• 组合生成代替逐一枚举：数字的幂和只和各位数字的出现次数有关，和顺序无关，直接生成所有不重复的数字组合计算幂和，再校验幂和的数字组成是否匹配即可，比遍历所有数的效率高几个数量级
• 预计算幂次表：提前把0-9的1~60次幂全部算好存起来，不用每次计算数字幂次，直接查表即可
• 替换素性判断算法：用确定性米勒-拉宾素性测试代替试除法，对于2^64以内的数，用固定测试基就能100%准确判断素数，大数场景下速度比试除法快上千倍
• 额外剪枝：素数除了2、5之外末位只能是1、3、7、9，且除了2之外全是奇数，生成组合时可以直接过滤掉不符合要求的情况

优化后代码 #

# 预计算0-9的1~60次幂表
max_digits = 60
pow_table = [[0]*(max_digits+1) for _ in range(10)]
for d in range(10):
    for p in range(1, max_digits+1):
        pow_table[d][p] = d ** p

# 2^64以内确定性米勒-拉宾素性测试
def is_prime(n):
    if n < 2:
        return False
    # 小素数直接判断
    small_primes = [2,3,5,7,11,13,17,19,23,29,31,37]
    for p in small_primes:
        if n % p == 0:
            return n == p
    # 分解n-1为d*2^s
    d = n - 1
    s = 0
    while d % 2 == 0:
        d //= 2
        s += 1
    # 2^64以内固定测试基
    for a in [2, 325, 9375, 28178, 450775, 9780504, 1795265022]:
        if a % n == 0:
            continue
        x = pow(a, d, n)
        if x == 1 or x == n-1:
            continue
        for _ in range(s-1):
            x = pow(x, 2, n)
            if x == n-1:
                break
        else:
            return False
    return True

from itertools import combinations_with_replacement

def find_prime_narcissistic(target_count):
    res = []
    # 1位数单独处理
    for n in [2,3,5,7]:
        res.append(n)
        if len(res) == target_count:
            return res
    # 从2位到max_digits位遍历
    for digit_len in range(2, max_digits+1):
        # 生成升序的数字组合，避免重复计算
        for digits in combinations_with_replacement(range(10), digit_len):
            # 计算当前组合的幂和
            total = sum(pow_table[d][digit_len] for d in digits)
            # 幂和位数不对直接跳过
            if len(str(total)) != digit_len:
                continue
            # 校验幂和的数字排序后和原组合一致
            if tuple(sorted(int(c) for c in str(total))) == digits:
                # 判断是不是素数
                if is_prime(total):
                    res.append(total)
                    if len(res) == target_count:
                        return res
    return res

# 找7个素数自幂数
result = find_prime_narcissistic(7)
for num in result:
    print(num)

运行说明 #

这套代码几秒内就能跑完得到全部7个素数自幂数，分别是2、3、5、7、28116440335967、4498128791164624869、4338281769391371。

内容的提问来源于stack exchange，提问作者wallshock