Foobar挑战质数位置任务遇阻：大输入超时问题求解

看起来你在Foobar的质数位置任务里遇到了性能瓶颈，尤其是当N=10000时的超时问题。我来帮你分析一下可能的低效点，然后给出优化后的解决方案。

问题核心分析 #

你的代码超时大概率是两个原因：

• 字符串拼接效率低下：如果你的代码是用+=来逐个拼接质数字符串，那随着字符串长度增加，每次拼接都会复制整个现有字符串，时间复杂度会变成O(k²)（k是总字符数），当k达到20000+时，这个开销会非常大。
• 质数生成或判断不够高效：如果你的质数生成方法耗时，或者生成了过多不必要的质数，也会拖慢整体速度。

优化方案 #

我们可以从两个关键点入手优化：

• 用列表存储质数字符串片段，最后一次性join，把字符串拼接的时间复杂度降到O(k)。
• 用按需生成质数的生成器，只生成到我们需要的总字符长度为止，避免多余计算。

优化后的代码 #

def is_prime(num):
    if num < 2:
        return False
    if num == 2:
        return True
    if num % 2 == 0:
        return False
    # 只检查奇数除数，减少一半计算量
    for i in range(3, int(num**0.5) + 1, 2):
        if num % i == 0:
            return False
    return True

def prime_generator():
    num = 2
    while True:
        if is_prime(num):
            yield num
        num += 1

def answer(n):
    n = int(n)
    # 计算需要的最小总字符长度：起始位置n + 要取的n+4个字符
    required_total_length = n + (n + 4)
    prime_str_parts = []
    current_total_length = 0
    prime_gen = prime_generator()
    
    # 生成质数，直到总字符长度满足要求
    while current_total_length < required_total_length:
        prime = next(prime_gen)
        prime_str = str(prime)
        prime_str_parts.append(prime_str)
        current_total_length += len(prime_str)
    
    # 拼接成完整的质数字符串
    full_prime_str = ''.join(prime_str_parts)
    # 截取目标子串
    return full_prime_str[n:n + n + 4]

代码解释 #

1. 质数判断优化：is_prime函数跳过了偶数除数，只检查从3开始的奇数，减少了一半的循环次数，判断小质数的速度更快。
2. 质数生成器：prime_generator是一个无限生成器，只会在需要的时候生成下一个质数，不会提前生成过多质数浪费资源。
3. 高效字符串拼接：用列表prime_str_parts存储每个质数的字符串，最后用''.join()拼接，这比反复用+=高效得多，因为列表的append操作是O(1)的，而join是一次性线性操作。
4. 按需停止生成：我们计算出需要的最小总字符长度，一旦当前拼接的总长度达到这个值，就停止生成质数，避免做无用功。

测试验证 #

• 当n=0时，代码会生成质数直到总长度≥5，得到"235711"，截取[0:5]得到"23571"，符合要求。
• 当n=3时，总长度需要≥8，生成的字符串是"23571113"，截取[3:8]得到"71113"，正确。

这个版本的代码在N=10000时应该能快速运行，因为所有操作都是线性或接近线性的时间复杂度，没有多余的开销。

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