为何优化后的素数生成器程序耗时反而更长？

我最近在开发素数生成器，尝试优化时间复杂度：

• 初始实现：从current_big_prime + 1开始迭代current_number，检查素数时遍历已找到的所有素数。
• 优化思路：理论上检查素数时只需遍历到sqrt(current_number)+1，应该能大幅节省时间，但实际运行后，优化后的程序性能反而更差。

以下是我的测试代码：

#%% Prime Streaming (PG-13)
import timeit
import matplotlib.pyplot as plt
import math

class Primes:
    @staticmethod
    def stream1():
        primes = []
        yield 2
        i = 3
        while True:
            is_prime = True
            for j in primes:
                if i % j == 0: 
                    is_prime = False
                    break
            if is_prime:
                yield i
                primes.append(i)
            i += 2
    @staticmethod
    def stream2():
        primes = []
        yield 2
        i = 3
        while True:
            is_prime = True
            for j in filter(lambda x: x < math.sqrt(i) + 1, primes):
                if i % j == 0: 
                    is_prime = False
                    break
            if is_prime:
                yield i
                primes.append(i)
            i += 2

def get_nth_prime1(n):
    s = Primes.stream1()
    cnt = 0
    for i in s:
        cnt += 1
        if cnt > n:
            break
    return i
def get_nth_prime2(n):
    s = Primes.stream2()
    cnt = 0
    for i in s:
        cnt += 1
        if cnt > n:
            break
    return i

s1 = Primes().stream1()
for i in range(20):
    print(next(s1), end=' ')
print()

s2 = Primes().stream2()
for i in range(20):
    print(next(s2), end=' ')
print()

t1 = timeit.timeit("get_nth_prime1(5000)",  
              "from __main__ import get_nth_prime1", number=1)
t2 = timeit.timeit("get_nth_prime2(5000)",  
              "from __main__ import get_nth_prime2", number=1)
print(f"check until n: {t1}\ncheck until sqrt(n): {t2}")

性能下降的原因 #

stream2的问题出在filter(lambda x: x < math.sqrt(i) + 1, primes)这个操作上：

1. 无意义的全量判断：filter会对primes列表里的每个元素都执行一次lambda判断，哪怕某个元素已经超过sqrt(i)+1，还是会继续判断后面的元素，无法提前终止遍历。而stream1虽然遍历全部素数，但一旦找到能整除的就break，实际开销反而更小。
2. 额外的函数调用开销：lambda表达式和filter函数本身存在调用成本，这种开销抵消了减少遍历次数带来的性能收益，甚至超过。

正确的优化实现 #

应该提前计算当前数的平方根，遍历素数时一旦素数超过这个值就直接终止循环，既减少遍历次数，又避免额外的函数开销：

@staticmethod
def stream3():
    primes = []
    yield 2
    i = 3
    while True:
        is_prime = True
        sqrt_i = math.isqrt(i)  # 使用整数平方根，比sqrt更高效且避免浮点误差
        for j in primes:
            if j > sqrt_i:
                break  # 超过平方根后无需继续判断
            if i % j == 0:
                is_prime = False
                break
        if is_prime:
            yield i
            primes.append(i)
        i += 2

这个版本的性能会明显优于前两个，因为它精准地控制了遍历的终止时机，没有多余的计算开销。

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