如何在Numpy中基于元素索引实现条件索引修改数组？

核心思路是利用Numpy的广播机制和布尔索引，完全规避Python层面的循环枚举——这是Numpy处理大规模数组的核心优化手段，效率比手写循环高几个数量级。

根据scale数组的维度不同，分三种常见场景实现：

场景1：scale是与二维数组行对应的一维数组 #

假设你的二维数组array形状为(n_rows, n_cols)，scale是长度为n_rows的一维数组（每个元素对应array的一行），需要先将scale调整为(n_rows, 1)的形状，让它能和array的每一列广播相乘：

import numpy as np

# 示例数据
array = np.random.rand(5, 3)  # 5行3列的二维数组
scale = np.random.rand(5)    # 对应每一行的缩放因子
threshold = 0.6

# 生成布尔掩码：每个元素乘以对应行的scale后大于阈值
mask = (array * scale.reshape(-1, 1)) > threshold
# 等价写法：scale[:, np.newaxis] 替代 scale.reshape(-1,1)

# 对满足条件的元素赋值为np.nan
array[mask] = np.nan

场景2：scale是与二维数组列对应的一维数组 #

如果scale是长度为n_cols的一维数组（每个元素对应array的一列），Numpy会自动将其广播到每一行，无需额外调整形状：

array = np.random.rand(5, 3)
scale = np.random.rand(3)    # 对应每一列的缩放因子
threshold = 0.6

mask = (array * scale) > threshold
array[mask] = np.nan

场景3：scale与二维数组同形状 #

如果scale本身就是和array完全同形状的二维数组，直接相乘生成掩码即可：

array = np.random.rand(5, 3)
scale = np.random.rand(5, 3)  # 和array同形状的缩放数组
threshold = 0.6

mask = (array * scale) > threshold
array[mask] = np.nan

关键说明 #

• Numpy的向量操作是底层用C实现的，完全避开了Python循环的开销，数组规模越大，效率优势越明显。
• 广播规则：两个数组从后往前匹配维度，要么维度长度相同，要么其中一个维度长度为1，就能自动完成广播。如果不确定形状是否兼容，可以用np.expand_dims或reshape手动调整。

内容的提问来源于stack exchange，提问作者E. Gptreas