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

【问题标题】：Python: how to get all the first values row-wise from a 2D numpy array when using a 2D boolean maskPython：如何在使用 2D 布尔掩码时从 2D numpy 数组中逐行获取所有第一个值
【发布时间】：2021-12-02 11:06:55
【问题描述】：

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

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

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

【问题讨论】：

标签： python arrays numpy performance vectorization

【解决方案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()

【讨论】：