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()