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!
CodePudding user response:
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()