如何基于行阈值数组，将NumPy矩阵列索引超阈值的元素设为0

实现方案 #

你可以通过NumPy的广播机制实现全量的索引判断，不需要逐行遍历，效率很高，完整代码如下：

import numpy as np

# 定义阈值数组与原始矩阵
thresh_arr = np.array([0, 0, 2, 1])
raw_matrix = np.ones((4, 3), dtype=int)

# 生成所有列的索引数组
col_index = np.arange(raw_matrix.shape[1])
# 扩展阈值数组维度后与列索引比较，生成符合要求的布尔掩码
mask = col_index > thresh_arr[:, np.newaxis]
# 将掩码对应位置的元素赋值为0
raw_matrix[mask] = 0

print(raw_matrix)

运行后输出结果和你预期完全一致：

[[1 0 0]
 [1 0 0]
 [1 1 1]
 [1 1 0]]

原理解释 #

1. 整个逻辑完全基于列索引和行阈值的比较，和矩阵本身的元素值无关，符合你的要求。
2. thresh_arr[:, np.newaxis]的作用是把1维的阈值数组转为形状为(4,1)的列向量，和形状为(3,)的列索引数组做比较时，NumPy会自动将两个数组广播为和原始矩阵同形状的(4,3)布尔矩阵，每个位置的布尔值就代表当前位置是否需要改为0。
3. 如果觉得np.newaxis不好记，也可以替换为thresh_arr.reshape(-1, 1)，效果完全相同。

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