ndarray построчный индекс значений больше массива

Question

У меня есть один массив формы (X, 5):

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

и один массив формы (X, 1):

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

Теперь мне нужно получить первый индекс M> = n для каждой строки или nan, если такого индекса нет. Например:

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

Эти примеры применяются к каждой строке отдельно (я мог бы зациклить X раз, так как оба массива имеют длину X).

Интересно, есть ли способ сделать это многомерным способом, чтобы получить вывод всех индексов одновременно?

Что-то вроде:

np.where(M>=n)

Спасибо

Редактировать: значения в M не отсортированы, я все еще ищу первый подходящий индекс / вхождение M> = n (поэтому, вероятно, не отсортированный)

Answer 1

Вы можете начать с проверки, какие индексы строк меньше или равны n, и использовать argmax, чтобы получить первые True для каждой строки.Для строк, где все столбцы False, мы можем использовать np.where, чтобы установить их на np.nan, например:

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