如何基于第一个数组非零条件实现NumPy数组间的相减操作

问题原因分析 #

• 第一种掩码赋值报错的原因是：a[a!=0] 提取出的是所有非零元素组成的一维数组，而你直接把整个二维数组b赋值给它，维度不匹配导致报错。
• 第二种np.subtract的where参数出现随机值的原因是：where仅会在条件为真的位置执行计算，未匹配位置不会被初始化，会保留内存中的随机垃圾值。

可行实现方案 #

以下方案均为纯向量化操作，无遍历，性能和NumPy原生运算一致，适合大规模数组：

方案1：修正掩码赋值（最直观，易理解）

import numpy as np
a = np.array([[1, 2, 3], [4, 0, 0], [7, 8, 0]])
b = np.array([[10, 11, 12], [13, 14, 15], [16, 17, 18]])

mask = a != 0
res = np.zeros_like(a)
res[mask] = a[mask] - b[mask]

如果允许修改原数组a，可以直接写：

a[mask] = a[mask] - b[mask]

方案2：修正np.subtract用法（一行实现）

给np.subtract指定out参数初始为全0数组，未匹配位置会保留初始的0值：

res = np.subtract(a, b, where=a!=0, out=np.zeros_like(a))

方案3：布尔掩码乘法（最简洁）

利用布尔数组运算时会自动转为1/0的特性，直接对差值做掩码过滤：

res = (a - b) * (a != 0)

三种方案输出均符合预期：

array([[-9, -9, -9],
       [-9,  0,  0],
       [-9, -9,  0]])

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