Python打印n个素数传参100时返回101个素数的bug排查

问题代码 #

def prime(n):
    if n==2 or n==3:
        return True
    if n==1:
        return False
    f=True
    for i in range(2,int(n**0.5)+1):
        if n%i==0:
            f=False 
            break
    return f
def n_primes(n):
    primes=[2,3,5,7]
    k=10
    count=4
    itr=[1,3,7,9]
    while count!=n:
        for i in itr:
            if prime(k+i):
               primes.append(k+i)
               count+=1 
        k+=10
    return primes
print(len(n_primes(100)))

问题现象 #

传入n=100调用函数时，实际返回长度为101的素数列表；更换循环终止判断条件（用素数列表长度替代count变量计数）后，异常仍然存在。

根因定位 #

逻辑漏洞出在两个位置：

• 循环终止判断粒度过粗：while循环仅在每跑完一整轮10区间内的4个候选数（尾数为1、3、7、9的数）后，才会检查计数是否达到目标值n。遍历4个候选数的for循环内部没有计数达标即终止的逻辑，只要进入某一轮10区间的遍历，就会把该区间内所有符合条件的素数全部加入列表，哪怕计数在遍历中途就已经达到n值。
• 终止条件存在风险：用count!=n做判断，一旦count因为批量加元素超过n，循环就不会自动停止，存在死循环隐患。

n=100时多返回1个素数的直接原因是：第100个素数541和第101个素数547刚好落在同一个10区间（k=540段）内，遍历到541时count已经达到100，但循环没有停止，继续检查到同区间的547也是素数，就多加入了这一个元素。

注：之前换用列表长度做判断没有效果，本质是判断逻辑还是放在while层，没有深入到for循环内部，批量加元素的逻辑没有发生变化。

修复方案 #

在for循环内每次新增素数、计数自增后，立刻判断是否达到目标数量，达标则终止当前轮次的遍历，同时把while循环的判断条件改为count < n避免死循环：

def n_primes(n):
    # 兼容n小于初始素数列表长度的边界场景
    if n <= 4:
        return [2,3,5,7][:n]
    primes=[2,3,5,7]
    k=10
    count=4
    itr=[1,3,7,9]
    while count < n:
        for i in itr:
            num = k + i
            if prime(num):
                primes.append(num)
                count += 1
                # 计数达标立刻跳出当前for循环
                if count == n:
                    break
        k += 10
    return primes

修复后传入n=100时，函数会准确返回100个素数，不会出现多算、死循环的问题。

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