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

Question

I have two large 2D arrays, one with values and the other ones with a mask of "valid" values.

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

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

My goal is to get, for each row, the first value when valid==True , and obtain a vector of that sort: [2, 9, 3] , in the fastest possible way.

I tried applying the mask and querying from it, but it destroys the structure:

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

I tried looping through all the indices, but I am wondering if there is a faster and vectorized way of doing that. Thank you!

Answer 1

Try:

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

Or use np.take_along_axis :

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