Python NumPy数组条件替换技巧

np.where(condition, x, y)中condition必须为布尔数组，x和y可广播，返回新数组；多条件优先用np.select，避免嵌套；注意维度对齐与dtype推断。

NumPy where 函数怎么写对：别把条件和替换值顺序搞反

np.where 的签名是 np.where(condition, x, y)，不是 np.where(condition, true_value, false_value) —— 这俩看着像，但 x 和 y 是数组（或标量），会广播参与计算。很多人卡在第一步：传了三个标量以为能“原地替换”，结果返回的是新数组，且形状由广播规则决定。

常见错误现象：np.where(arr > 5, 99, arr) 想把大于 5 的元素换成 99，却误写成 np.where(arr > 5, 99, 0)，导致全数组变成 99 或 0，丢失原值。

• 条件必须是布尔数组，arr > 5 可以，arr[arr > 5] 不行（那是索引，不是条件）
• x 和 y 可以是标量、同形数组、或可广播的数组（比如 x=99 是标量，y=arr 是原数组，合法）
• 返回值永远是新数组，不会修改原 arr；如需就地修改，得赋值回原变量：arr = np.where(...)

用 where 做多条件替换：嵌套 or 链式？选后者更清晰

两个以上分支时，嵌套 np.where（比如 np.where(a>10, 1, np.where(a>5, 2, 0))）语法合法，但可读性差、易出括号错。更推荐用布尔组合 + 逐层覆盖，或直接用 np.select。

例如：将数组按区间映射为类别码（0–5→0，6–10→1，11+→2）：

cond1 = arr <= 5
cond2 = (arr > 5) & (arr <= 10)  # 注意用 &，不是 and
cond3 = arr > 10
result = np.select([cond1, cond2, cond3], [0, 1, 2], default=-1)

• np.select 的条件列表必须互斥（或按顺序匹配第一个为 True 的），不互斥时只取首个命中项
• 用 & / | / ~ 拼布尔条件，绝不用 and/or（会报 ValueError: The truth value of an array with more than one element is ambiguous）
• 如果坚持用嵌套 where，建议拆成变量分步写，避免一行嵌三层

性能陷阱：where 在大数组上广播 x 或 y 会吃内存

当 x 或 y 是跟原数组同 shape 的数组时，np.where 内部要复制/广播它们参与三元运算——这会瞬间多占一倍内存。例如：np.where(mask, big_arr * 2, big_arr) 中，big_arr * 2 会新建一个同样大小的临时数组。

• 若只是替换少量值，优先用布尔索引直接赋值：arr[mask] = 99（就地、无额外内存）
• 若必须用 where 做复杂逻辑，且 x/y 是计算结果，考虑用 np.copyto + out 参数控制输出位置
• 注意 dtype 自动推断：若 x 是 int、y 是 float，结果会升为 float64，可能影响后续计算精度或内存占用

为什么 where 返回值有时 shape 不对？检查布尔条件是否被误广播

典型问题：arr 是 (100, 3) 数组，mask = arr[:, 0] > 5 得到 (100,) 布尔向量，但 np.where(mask, val1, val2) 报错或返回奇怪 shape —— 因为 mask 和 arr 维度不匹配，where 尝试广播失败。

• 确保 condition 的 shape 能与 x、y 广播对齐；最稳妥是让 condition 和 x/y 同 shape
• 用 np.expand_dims 或切片加 None 对齐维度，比如 mask[:, None] 把 (100,) 变成 (100, 1)，就能和 (100, 3) 的 arr 广播
• 调试时先打印 mask.shape、x.shape、y.shape，比猜快得多

真正麻烦的从来不是语法记不住，而是布尔条件维度没对齐，或者以为 where 会就地改原数组——它不会，连 dtype 都可能悄悄变。