Python Numpy：同形状多数组需相同修改时，有无更快实现方式？

你有多个形状完全相同的NumPy数组，以及两个同形状的掩码，目前通过逐个修改每个数组的方式处理，想要更高效的实现方案。原代码示例如下：

import time
import numpy as np

start = time.time()

mask1 = np.array([[1,0,1],[0,1,1],[1,0,1]])
mask2 = np.array([[0,0,0],[1,0,1],[0,0,1]])

arr1 = np.array([[20,10,51],[21,1,2],[25,23,38]])
arr2 = np.array([[99,1,6],[66,54,11],[22,21,1]])
arr3 = np.array([[23,2,3],[55,2,16],[90,37,1]])
arr4 = np.array([[81,25,22],[1,63,24],[47,58,1]])

arr1[(mask1 == 1) & (mask2 == 0)] = 9999
arr2[(mask1 == 1) & (mask2 == 0)] = 9999
arr3[(mask1 == 1) & (mask2 == 0)] = 9999
arr4[(mask1 == 1) & (mask2 == 0)] = 9999

print(time.time() - start)

优化方案 #

核心优化点

• 预计算掩码条件：避免每次赋值都重复计算(mask1 == 1) & (mask2 == 0)，减少冗余运算
• 批量处理数组：利用NumPy的向量化操作，将多个同形状数组堆叠后一次性修改，大幅提升效率

优化代码示例

import time
import numpy as np

start = time.time()

mask1 = np.array([[1,0,1],[0,1,1],[1,0,1]])
mask2 = np.array([[0,0,0],[1,0,1],[0,0,1]])

arr1 = np.array([[20,10,51],[21,1,2],[25,23,38]])
arr2 = np.array([[99,1,6],[66,54,11],[22,21,1]])
arr3 = np.array([[23,2,3],[55,2,16],[90,37,1]])
arr4 = np.array([[81,25,22],[1,63,24],[47,58,1]])

# 提前计算合并后的掩码，仅执行一次
combined_mask = (mask1 == 1) & (mask2 == 0)

# 将所有数组堆叠成一个高维数组，一次性完成赋值
stacked_arrays = np.stack([arr1, arr2, arr3, arr4])
stacked_arrays[:, combined_mask] = 9999

# 如需还原为单独数组，拆分即可
arr1, arr2, arr3, arr4 = stacked_arrays

print(time.time() - start)

备选方案（适合不想堆叠数组的场景）

如果需要保留原数组的独立结构，也可以用列表存储数组，配合预计算的掩码循环赋值，同样比原方法高效：

import time
import numpy as np

start = time.time()

mask1 = np.array([[1,0,1],[0,1,1],[1,0,1]])
mask2 = np.array([[0,0,0],[1,0,1],[0,0,1]])

arr1 = np.array([[20,10,51],[21,1,2],[25,23,38]])
arr2 = np.array([[99,1,6],[66,54,11],[22,21,1]])
arr3 = np.array([[23,2,3],[55,2,16],[90,37,1]])
arr4 = np.array([[81,25,22],[1,63,24],[47,58,1]])

combined_mask = (mask1 == 1) & (mask2 == 0)
# 用列表存储数组，循环赋值
for arr in [arr1, arr2, arr3, arr4]:
    arr[combined_mask] = 9999

print(time.time() - start)

效果说明 #

两种优化方案都避免了重复计算掩码，其中堆叠数组的方式利用了NumPy底层的向量化优化，在数组数量越多时，效率提升越明显。

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