如何基于1D数组范围对Numpy二维数组进行行索引掩码？

给定1D数组ranges（每个元素对应二维数组的列有效行数）和二维数组matrixes，要生成每列前ranges[i]行为True、其余为False的掩码数组，可通过NumPy的广播特性实现：

实现步骤 #

1. 获取二维数组的行列数，确定行索引范围
2. 将行索引转换为列向量，利用广播机制与ranges进行元素级比较

完整代码 #

import numpy as np

ranges = np.array([1, 2, 0])
matrixes = np.array([[0, 1, 2],
                     [3, 4, 5],
                     [6, 7, 8]])

# 获取行列数
rows, cols = matrixes.shape
# 生成列向量形式的行索引，形状为(rows, 1)
row_indices = np.arange(rows)[:, np.newaxis]
# 广播比较，生成掩码
mask = row_indices < ranges

print(mask)

输出结果 #

[[ True  True False]
 [False  True False]
 [False False False]]

原理说明 #

row_indices被转换为列向量后，会和ranges自动广播为相同形状的二维数组。每个位置的比较逻辑为：当前行索引是否小于对应列的有效行数阈值，直接得到符合要求的掩码矩阵。

内容的提问来源于stack exchange，提问作者Andrei M.