如何为二维NumPy数组按行生成指定列范围的布尔掩码？

问题描述 #

给定二维NumPy数组：

import numpy as np
data = np.ones((3, 5))

以及存储每行结束索引的一维数组：

end_indices = np.array([0, 2, 3])

需要生成布尔掩码，让每行中结束索引及之前的元素为True，期望输出：

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

解决方案：利用NumPy广播特性 #

这是最简洁高效的实现方式，无需循环，完全基于NumPy的向量化操作：

import numpy as np

data = np.ones((3, 5))
end_indices = np.array([0, 2, 3])

# 生成与data列数匹配的索引数组，并扩展为行向量
col_indices = np.arange(data.shape[1])[np.newaxis, :]
# 将结束索引扩展为列向量，与列索引数组广播后逐元素比较
mask = col_indices <= end_indices[:, np.newaxis]

print(mask)

原理说明 #

1. np.arange(data.shape[1])生成列索引数组[0, 1, 2, 3, 4]，通过[np.newaxis, :]将其转换为形状为(1, 5)的二维行向量；
2. end_indices[:, np.newaxis]将一维的结束索引数组转换为形状为(3, 1)的二维列向量；
3. 利用NumPy的广播机制，行向量和列向量会自动扩展为(3, 5)的二维数组，逐元素执行<=比较，直接得到目标布尔掩码。

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