найти индекс наибольшего, второго по величине или третьего по величине в ряду многомерного массива Пихтона

Question

Допустим, у меня есть следующие два пустых массива:

a = numpy.array([[1,4,6,2,5],[3,2,7,12,1],[8,5,3,1,4],[6,10,2,4,9]])
b = numpy.array([0, 1, 4])

Теперь я хочу сначала найти индекс максимального значения в определенной строке (просто произнесите вторую строку a[1,:]),У меня есть другой массив b с некоторыми числами в нем, и если индекс максимального значения присутствует в качестве элемента в другом b, я использую этот индекс для дальнейшего вычисления.Если индекс максимального значения из строки a отсутствует как элемент в b, мне нужно искать индекс 2-го по величине числа, и если это значение индекса присутствует в b какэлемент, я иду на это в противном случае искать индекс третьего по величине числа и так далее.Я не хочу сортировать массив a.

В приведенном выше примере у меня есть [3,2,7,12,1] во второй строке. Индекс максимального числа равен 3, но этого 3 нет в b, тогда индекс второго по величине равен 2, что также неприсутствует в б, то я ищу индекс третьего по величине, который равен 0, и это присутствует в б.Поэтому я назначаю 0 новой переменной.Какой-нибудь быстрый и быстрый способ сделать это?Заранее спасибо.

Answer 1

Вот тот, который хорошо масштабируется для общих ndarrays -

def maxindex(a, b, fillna=-1):
    sidx = a.argsort(-1)
    m = np.isin(sidx,b)
    idx = m.shape[-1] - m[...,::-1].argmax(-1) - 1
    out = np.take_along_axis(sidx,idx[...,None],axis=-1).squeeze()
    return np.where(m.any(-1), out, fillna)

Образцы прогонов -

In [83]: a
Out[83]: 
array([[ 1,  4,  6,  2,  5],
       [ 3,  2,  7, 12,  1],
       [ 8,  5,  3,  1,  4],
       [ 6, 10,  2,  4,  9]])

In [84]: b
Out[84]: array([0, 1, 4])

In [85]: maxindex(a, b) # all rows
Out[85]: array([4, 0, 0, 1])

In [86]: maxindex(a[1], b) # second row
Out[86]: array([0])

3D-кейс -

In [105]: a
Out[105]: 
array([[[ 1,  4,  6,  2,  5],
        [ 3,  2,  7, 12,  1],
        [ 8,  5,  3,  1,  4],
        [ 6, 10,  2,  4,  9]],

       [[ 1,  4,  6,  2,  5],
        [ 3,  2,  7, 12,  1],
        [ 8,  5,  3,  1,  4],
        [ 6, 10,  2,  4,  9]]])

In [106]: maxindex(a, b)
Out[106]: 
array([[4, 0, 0, 1],
       [4, 0, 0, 1]])

Answer 2

IIUC, если вам нужно сохранить индекс от a:

res = b[a[:, b].argmax(1)]

array([4, 0, 0, 1])

В качестве альтернативы:

a[:, np.delete(np.arange(a.shape[1]), b)] = a.min()
res = a.argmax(1)

array([4, 0, 0, 1], dtype=int64)

Если индекса массивного массива достаточно:

res = a[:, b].argmax(1)

array([2, 0, 0, 1], dtype=int64)