如何使用Numpy以向量化方式按行生成二维数组的指定索引上限掩码？

给定如下Numpy数组：

>>> import numpy as np
>>> A = np.arange(2*3).reshape(2,3)
array([[0, 1, 2],
       [3, 4, 5]])

>>> b = np.array([1, 2])

需要生成每行前b[i]个元素为True的掩码：

>>> mask
array([[True, False, False],
       [True, True, False]])

完全可以用Numpy的向量化操作实现，核心利用广播机制让每行的列索引与对应行的上限值做比较，无需循环。

实现代码 #

# 生成与A列数匹配的列索引数组
col_indices = np.arange(A.shape[1])
# 通过广播实现逐行比较，生成掩码
mask = col_indices < b[:, np.newaxis]

原理说明 #

• col_indices = np.arange(A.shape[1])生成数组[0,1,2]，对应A的列索引；
• b[:, np.newaxis]将一维数组b转换为二维列数组[[1],[2]]，触发Numpy广播机制，让列索引数组与每行的上限值逐一比较；
• 比较结果直接输出符合要求的布尔掩码。

验证结果 #

>>> mask
array([[ True, False, False],
       [ True,  True, False]])

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