Python：使用 2D 布尔掩码时如何从 2D numpy 数组逐行获取所有第一个值

Question

我有两个大型二维数组，一个有值，另一个有“有效”值掩码。

vals = np.array([
    [5, 2, 4],
    [7, 8, 9],
    [1, 3, 2],
])

valid = np.array([
    [False, True, True],
    [False, False, True],
    [False, True, True],
])

我的目标是在 valid==True 时为每一行获取第一个值，并以最快的方式获取该类型的向量：[2, 9, 3]。

我尝试应用掩码并从中查询，但它破坏了结构：

vals[valid]
> array([2, 4, 9, 3, 2])

我尝试遍历所有索引，但我想知道是否有更快的矢量化方法来执行此操作。谢谢！

Answer 1

尝试：

vals[np.arange(len(vals)), np.argmax(valid,axis=1)]

或使用np.take_along_axis:

np.take_along_axis(vals, np.argmax(valid,axis=1)[:,None], axis=1).ravel()

Python：使用 2D 布尔掩码时如何从 2D numpy 数组逐行获取所有第一个值

Python: how to get all the first values row-wise from a 2D numpy array when using a 2D boolean mask

python

arrays

performance

numpy

vectorization