如何加速含np.sum的循环代码？Python数值计算优化求助

首先，针对你这段要执行1000万+次的代码，核心优化思路是减少重复计算、降低Python循环的开销——毕竟Python的for循环在高频场景下本身就偏慢，而且你原来的代码里对每个block重复计算了两次np.sum，这完全是性能浪费。

基础优化：消除重复计算与批量汇总 #

先看你原来的代码：

mi = np.random.randint(2,size=(4,4))
list_of_blocks = [np.array([[0],[1]]),np.array([[0,0,0],[0,1,0]])]
nee = np.sum(mi)
n1 =0 
mrxnr= 0
for block in list_of_blocks:
    nee -= np.sum(block)
    n1 += np.sum(block)
    mrxnr += block.shape[0]*block.shape[1]

这里最大的问题是：每个block的np.sum被计算了两次，且循环内做了三次独立操作。我们可以把所有block的总和、元素总数提前批量计算，再一次性完成变量更新：

mi = np.random.randint(2,size=(4,4))
list_of_blocks = [np.array([[0],[1]]),np.array([[0,0,0],[0,1,0]])]
nee = np.sum(mi)

# 只遍历一次list_of_blocks，批量计算核心汇总值
total_block_sum = sum(np.sum(b) for b in list_of_blocks)
total_elements = sum(b.size for b in list_of_blocks)

# 一次性完成变量更新，完全消除循环内的冗余操作
nee -= total_block_sum
n1 = total_block_sum
mrxnr = total_elements

优化点解释：

1. 减少重复计算：每个block的总和只计算一次，砍掉一半的np.sum调用开销；
2. 替换shape计算：用b.size代替block.shape[0]*block.shape[1]，size是numpy数组的内置属性，直接返回元素总数，比手动维度相乘更快；
3. 降低循环开销：原来的循环需要迭代每个block并执行三次操作，现在只需要一次遍历完成汇总，剩下的都是简单加减赋值，Python层面的循环开销大幅降低。

极致优化：提前计算固定值（如果list_of_blocks不变） #

如果你的list_of_blocks在整个仿真过程中是固定不变的，那可以把total_block_sum和total_elements的计算提到高频循环的外面，只执行一次——这样每次1000万+次的循环里完全不需要再处理list_of_blocks：

# 初始化阶段：只执行一次
list_of_blocks = [np.array([[0],[1]]),np.array([[0,0,0],[0,1,0]])]
total_block_sum = sum(np.sum(b) for b in list_of_blocks)
total_elements = sum(b.size for b in list_of_blocks)

# 高频执行的循环（10mi+次）
for _ in range(10**7):
    mi = np.random.randint(2,size=(4,4))  # 假设mi每次都会变化
    nee = np.sum(mi) - total_block_sum
    n1 = total_block_sum
    mrxnr = total_elements

这种方式的性能提升最显著，因为完全消除了每次循环中遍历list_of_blocks的开销，只剩下最核心的np.sum(mi)和简单赋值。

终极优化：JIT编译（针对list_of_blocks动态变化的场景） #

如果list_of_blocks每次循环都不一样，没办法提前计算汇总值，那可以用Numba的JIT编译把代码转换成机器码，Python循环的速度会提升几十甚至上百倍：

from numba import jit

# 用numba编译函数，nopython=True强制生成纯机器码
@jit(nopython=True)
def compute_metrics(mi, list_of_blocks):
    nee = mi.sum()
    n1 = 0
    mrxnr = 0
    for block in list_of_blocks:
        block_sum = block.sum()
        nee -= block_sum
        n1 += block_sum
        mrxnr += block.size
    return nee, n1, mrxnr

# 使用示例
mi = np.random.randint(2,size=(4,4))
list_of_blocks = [np.array([[0],[1]]),np.array([[0,0,0],[0,1,0]])]
nee, n1, mrxnr = compute_metrics(mi, list_of_blocks)

Numba会把这个函数编译成机器码，完全绕过Python解释器的开销，循环执行效率会和C语言差不多，非常适合高频执行的场景。

额外小技巧 #

如果你的block里都是0和1这类整数，还可以尝试用np.count_nonzero(block)代替np.sum(block)——统计非零元素的速度有时候会比求和更快（尤其是当0占比很高的时候），可以根据你的实际数据测试一下哪个更高效。

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