如何无循环高效实现带条件的两个Numpy数组相减？

这事儿完全不用写循环，Numpy自带的np.where函数就能帮你搞定，而且是向量化运算，效率拉满！

核心思路 #

我们可以利用np.where的向量化条件判断特性：它会逐个元素检查条件，满足条件的位置赋值为指定值，不满足的位置保留另一个数组的对应元素，完美匹配你的需求：

• 条件：b > 0
• 满足条件时：赋值255
• 不满足条件时：取a的对应元素

完整代码示例 #

import numpy as np

# 定义原始数组
a = np.array([[[1,2,3], [4,5,6]], [[7,8,9], [10,11,12]]])
b = np.array([[[1,1,1], [0,0,0]], [[3,3,3], [4,4,4]]])

# 执行条件运算
c = np.where(b > 0, 255, a)

# 查看结果
print(c)

运行结果 #

输出正好是你预期的：

[[[255 255 255]
  [  4   5   6]]

 [[255 255 255]
  [255 255 255]]]

为什么不用循环？ #

Numpy的向量化操作是底层用C实现的，相比Python层面的循环，处理大规模数组时速度会快几个数量级，而且代码更简洁易读。

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