如何使用NumPy的np.where()函数实现条件替换操作？

本文共计743个文字，预计阅读时间需要3分钟。

pythonnp.where(condition, x, y)

看起来像是Python的三元操作符 `x if cond else y`，但它是一种向量化操作。不逐个元素进行条件判断，而是先计算所有 `x` 和 `y` 的值，然后根据 `condition` 的布尔值进行选择。这意味着，如果 `condition` 为 `False`，则输出所有 `x` 的值，否则输出所有 `y` 的值。

• 如果 x 是 1 / arr，而 arr 含 0，哪怕只打算用 y 分支，也会触发 ZeroDivisionError
• np.where 不跳过计算，只跳过赋值
• 真正想“惰性求值”，得用布尔索引分步写：result = np.copy(y); result[condition] = x[condition]

替换 NaN 或无穷大时，用 isnan / isinf 别直接写 == np.nan

np.nan == np.nan 永远是 False，所以 arr == np.nan 全是 False，用它做 condition 会完全失效。

• 正确做法是用 np.isnan(arr) 检测 NaN，np.isinf(arr) 检测 ±inf
• 若需同时处理 NaN 和 inf：np.isnan(arr) | np.isinf(arr)（注意用 |，不是 or）
• 示例：把所有异常值替换成 0

  cleaned = np.where(np.isnan(arr) | np.isinf(arr), 0, arr)

多维数组里 condition 形状不匹配就报 ValueError

np.where 要求 condition、x、y 三者可广播（broadcastable）。常见翻车点：

• condition 是 (3,)，x 是 (3, 4)，y 是标量 → 行不通，因为 (3,) 和 (3, 4) 广播后是 (3, 4)，但 y 标量能广播，问题在 condition 维度不够
• 更隐蔽的是：condition 是 (3, 1)，x 是 (3, 4)，这时能广播，但结果可能不符合直觉（按列广播）
• 安全做法：显式对齐维度，比如用 np.expand_dims(condition, axis=1) 或直接用 condition[:, None]

只想改部分值？别硬套 np.where，布尔索引更直白

如果只是“把满足条件的位置替换成固定值”，比如把负数全设为 0，用 np.where 是杀鸡用牛刀。

• 直接写：arr[arr < 0] = 0，语义清晰、内存原地修改、不生成新数组
• np.where(arr < 0, 0, arr) 会强制返回新数组，多占一倍内存
• 只有当 x 和 y 都是动态表达式（比如两个不同数组的对应元素），才真正需要 np.where

复杂点在于 condition 的广播规则和 x/y 的计算时机——这两处不细看文档，很容易写出看似对、实则慢或报错的代码。