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