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

Question

我有两个大的2D数组，一个具有值，另一个具有掩码"valid"值。

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