numpy 如何只对非零元素做运算而不创建掩码数组

最常用方式是用np.flatnonzero()获取非零索引后原地运算：先idx=np.flatnonzero(arr)，再arr[idx]直接赋值或运算，避免布尔掩码内存开销，且彻底跳过零值计算，规避未定义操作风险。

直接对非零元素做运算而不显式创建掩码数组，核心思路是利用 numpy 的高级索引（fancy indexing）配合原地操作或视图更新，避免额外布尔数组占用内存。关键在于：用 np.nonzero() 或 np.flatnonzero() 获取索引，再用这些索引定位并修改对应位置的值。

用 np.nonzero() 获取索引后原地更新

这是最常用、最轻量的方式——不生成布尔掩码，只生成整数索引元组，然后直接索引赋值：

• idx = np.nonzero(arr) 返回一个元组，如 (array([0,2]), array([1,3]))，适用于多维数组
• 对一维数组，更高效用 idx = np.flatnonzero(arr)，返回一维整数数组
• 之后可直接用 arr[idx] += 1、arr[idx] **= 2 等原地运算

示例：

>>> import numpy as np
>> a = np.array([0, 2, 0, -3, 0, 4])
>> idx = np.flatnonzero(a) # 得到 [1, 3, 5]
>> a[idx] = np.sqrt(np.abs(a[idx])) # 只对非零元取平方根（绝对值后）
>> a
array([0. , 1.41421356, 0. , 1.73205081, 0. , 2. ])

对二维数组按行/列处理非零元

若需保持维度结构（比如每行独立缩放），可结合 np.nonzero() 和广播逻辑：

• rows, cols = np.nonzero(arr) 分别得到行索引和列索引
• 用 arr[rows, cols] 提取所有非零值，做统一运算后，再赋回原位置
• 若需按行归一化非零元，可先用 np.bincount(rows, weights=arr[rows,cols]) 统计每行非零和，再用索引广播除法

避免临时数组：用 out= 参数就地计算

当运算涉及函数（如 np.log, np.exp），可先用 np.where 做条件选择，但注意：它仍会计算所有元素。真正“跳过”零的计算，只能靠索引分步：

• 先提取非零值：nz =

  

  arr[np.flatnonzero(arr)]
• 对 nz 做运算（此时无冗余计算）
• 用索引写回：arr.flat[np.flatnonzero(arr)] = nz_new
• 对大数组，arr.flat 是扁平迭代器，比 arr.ravel() 更省内存（不复制）

注意边界：零本身参与运算时的替代策略

如果目标是“对非零元做某运算，零保持不变”，且该运算在零处未定义（如 1/x），就不能用 np.where(arr != 0, 1/arr, 0)——因为 1/arr 仍会触发除零警告甚至错误。正确做法是：

• 先取索引：nz_idx = np.flatnonzero(arr)
• 单独计算：arr[nz_idx] = 1.0 / arr[nz_idx]
• 这样零位置完全不参与任何除法，彻底规避问题