numpy sum的out与where参数使用结果不符的问题排查

问题场景 #

现有两个数组：

import numpy as np
a = np.array([1, 4, 2])
b = np.array([1, 0, 2])

期望得到数组c：其中b等于0的位置初始值为-5，b不等于0的位置为a和b对应元素的和，预期结果为array([2, -5, 4])。

执行以下代码后：

c = np.sum([a, b], out=np.zeros_like(a)-5, where=b!=0, axis=0)

实际得到结果array([2, 0, 4])，不符合预期。

问题原因 #

你误解了np.sum中where参数的作用：where是用来控制哪些元素参与求和计算，而非控制out数组的初始值保留。当where条件为False时，sum会将该位置的计算结果设为0（求和的默认初始值），直接覆盖掉out数组中预先设置的-5。

解决方案 #

方案1：分步初始化+条件赋值 #

先将结果数组初始化为-5，再对满足条件的位置单独计算求和结果：

import numpy as np

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

# 初始化结果数组为-5
c = np.full_like(a, -5)
# 筛选出b不等于0的位置
mask = b != 0
# 对满足条件的位置计算a和b的和
c[mask] = np.sum([a, b], axis=0)[mask]

运行后得到array([2, -5, 4])，符合预期。

方案2：结合np.where直接构造结果 #

用np.where根据条件选择求和结果或初始值，更简洁：

c = np.where(b != 0, np.sum([a, b], axis=0), -5)

同样能得到预期结果。

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