为何numpy.array在修改内部数据时性能更慢？如何优化？

你遇到的这个情况其实挺典型的——很多人刚上手numpy的时候都会踩这个坑！先给你理清楚原因，再说说怎么优化。

首先看你写的两个版本，核心逻辑几乎一模一样：都是用Python的for循环逐个检查、修改数组/列表里的元素。但numpy的优势根本不是单个元素的操作，它的强项是批量的向量化运算——也就是把一堆操作丢给numpy底层的C代码去批量执行，避开Python循环的开销。而你现在的写法，反而让numpy的单个元素操作的额外开销暴露出来了：numpy数组的每个元素访问都要做类型校验、从C内存到Python对象的转换，这些步骤比纯Python list的元素访问要重得多，自然就变慢了。

先贴一下你原来的两个实现，方便对比：

原始list版本 #

import math
def PrimeTable(n:int) -> list[int]:
   nb2m1 = n//2 - 1
   l = [True]*nb2m1
   for i in range(1,(math.isqrt(n)+1)//2):
      if l[i-1]:
         for j in range(i*3,nb2m1,i*2+1):
            l[j] = False
   return [2] + [i for i,v in zip(range(3,n,2), l) if v]

原始numpy版本 #

import math, numpy
def PrimeTable2(n:int) -> list[int]:
   nb2m1 = n//2 - 1
   l = numpy.full(nb2m1, True)
   for i in range(1,(math.isqrt(n)+1)//2):
      if l[i-1]:
         for j in range(i*3,nb2m1,i*2+1):
            l[j] = False
   return [2] + [i for i,v in zip(range(3,n,2), l) if v]

你给出的性能对比图也能直观看到差异：

（其中y0是list版本的PrimeTable，y1是numpy版本的PrimeTable2）

优化方案：用numpy的向量化操作替代Python循环 #

要让numpy发挥优势，关键是把内层的Python for循环换成numpy的批量切片赋值操作，让C层去处理批量修改，避开Python循环的开销。修改后的代码如下：

import math, numpy
def PrimeTable2(n:int) -> list[int]:
    if n < 2:
        return []
    if n == 2:
        return [2]
    nb2m1 = n//2 - 1
    l = numpy.full(nb2m1, True)
    sqrt_n = math.isqrt(n)
    for i in range(1, (sqrt_n + 1)//2):
        if l[i-1]:
            # 把原来的内层for循环换成numpy切片批量赋值
            start_idx = i * 3
            step = i * 2 + 1
            l[start_idx::step] = False
    return [2] + [i for i,v in zip(range(3, n, 2), l) if v]

优化原理 #

原来的内层for j in range(...): l[j] = False是在Python层面逐个修改numpy数组元素，每次操作都有额外开销；而改成l[start_idx::step] = False后，numpy会直接在底层C代码里完成整个切片的批量赋值，没有Python循环的 overhead，这才能真正利用numpy的高性能优势——尤其是当n非常大的时候，这个优化的性能提升会非常明显。

简单来说：numpy怕的是“小步慢走”的Python单个元素操作，爱的是“大步快跑”的批量向量化运算～

备注：内容来源于stack exchange，提问作者Mr. W