查找下一个素数的Python代码运行超时，该如何优化

原代码核心问题 #

• 缩进错误：贴出的代码中num = n + 1缩进与外层循环不对齐，存在格式问题
• 无意义外层循环：外层for i in range(1, 500)完全未使用i变量，还限制了素数查找范围，若输入n过大将直接返回错误结果
• 逻辑存在隐性缺陷：内层循环的遍历范围range(2, num)在循环启动时就已固定，即便循环内修改了num的值，也不会改变当前内层循环的遍历上限，会产生大量无效运算
• 素数判断效率极低：内层循环遍历范围覆盖2到num的所有整数，存在大量冗余判断
• 未过滤偶数：除2外所有素数均为奇数，原代码未跳过偶数判断，白白多执行一倍的循环次数

优化后实现代码 #

import math

def nextPrime(n):
    # 最小素数为2，处理输入小于2的场景
    if n < 2:
        return 2
    # 初始待判断数直接跳到下一个奇数，跳过偶数
    num = n + 1 if (n + 1) % 2 == 1 else n + 2
    while True:
        is_prime = True
        # 素数判断仅需遍历到√num，且仅判断奇数除数
        for j in range(3, int(math.sqrt(num)) + 1, 2):
            if num % j == 0:
                is_prime = False
                break
        if is_prime:
            return num
        # 每次加2跳过偶数，减少判断次数
        num += 2

优化效果说明 #

• 移除无意义外层循环，改用while循环实现无范围限制的素数查找，不会出现输入过大返回错误的问题
• 提前过滤所有偶数：待判断数字、遍历除数均只保留奇数，直接减少75%的循环次数
• 素数判断上限从原有的num缩小到√num，数字越大优化效果越显著，比如判断10000是否为素数，原逻辑需遍历9998次，优化后仅需遍历49次
• 匹配到可整除的除数后立刻终止当前判断循环，避免无效遍历

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