ndarray row-wise index of values greater than array

Question

I have one array of shape (X, 5):

M = [[1,2,3,4,5],
     [6,7,8,9,1],
     [2,5,7,8,3]
     ...]

and one array of shape (X, 1):

n = [[3],
     [7],
     [100],
     ...]

Now I need to get the first index of M >= n for each row, or nan if there is no such index. For example:

np.where([1,2,3,4,5] >= 3)[0][0] # Returns 2
np.searchsorted(np.array([1,2,3,4,5]), 3) # Returns 2

These examples are applied to each row individually (I could loop X times as both arrays have the length X).

I wonder, is there a way to do it in a multidimensional way to get an output of all indices at once?

Something like:

np.where(M>=n)

Thank you

Edit: Values in M are unsorted, I'm still looking for the first index/occurrence fitting M >= n (so probably not searchsorted)

Answer 1

You could start by checking which row indices are lower or equal than n and use argmax to get the first True for each row. For the rows where all columns are False, we can use np.where to set them to np.nan for instance:

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

n = np.array([[3],[7],[100]])

le = n<=M

# array([[False, False,  True,  True,  True],
#        [False,  True,  True,  True, False],
#        [False, False, False, False, False]])

lea = le.argmax(1)
has_any = le[np.arange(len(le)), lea]
np.where(has_any, lea, np.nan)
# array([ 2.,  1., nan])