如何高效对比两个NumPy数组并按条件修改目标数组

你需要的逻辑不用写循环，直接用NumPy原生的向量化操作就能实现，性能比np.nditer写的Python层循环高几个数量级，数组长度越大优势越明显。

方案1：逐元素相乘（最简实现，适配b为0/1数组场景） #

你的规则本质是「b对应位置为1则保留a原值，为0则a置0」，和两个数组逐元素相乘的逻辑完全一致：

import numpy as np

a = np.array([1,2,3,4,5,6])
b = np.array([0,1,1,0,1,0])

a = a * b

运行后得到的a就是你期望的array([0,2,3,0,5,0])。这个操作完全在C层执行，没有Python循环的调度开销，速度极快。

方案2：布尔索引批量赋值（逻辑更直观，适配b存在非0/1值场景） #

如果你的b数组可能存在0、1之外的其他值，只需要严格按照「b值为1才保留，否则置0」的规则，可以直接用布尔索引选中需要修改的位置批量赋值：

a[b != 1] = 0

这行代码会直接定位所有b不等于1的位置，把a对应位置的元素统一设为0，同样是向量化实现，性能和逐元素乘法基本一致，代码逻辑和你的需求描述完全对应，可读性更强。

为什么原来的np.nditer实现慢？ #

np.nditer只是提供了数组元素的迭代访问接口，你在循环内部写的if判断、逐元素赋值都是在Python解释器层面逐行执行的，完全没有利用NumPy的向量化加速能力，处理十万级以上长度的数组时，性能差距会达到上百倍。

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