Блог

Предположим, у меня есть массив a формы (2, 2, 2):

 a = np.array([[[7, 9],
               [19, 18]],
              [[24, 5],
               [18, 11]]])
 

и массив b , максимальный a размер которого составляет : b=a.max(-1) (по строкам):

 b = np.array([[9, 19],
              [24, 18]])
 

Я хотел бы получить индекс элементов при b использовании индекса в сплющенном a виде , т. Е. a.reshape(-1) :

 array([ 7,  9, 19, 18, 24,  5, 18, 11])
 

Результатом должен быть массив одинаковой формы b с индексами b в сплющенном виде a :

 array([[1, 2],
       [4, 6]])
 

В основном это результат maxpool2d, когда return_indices= True в pytorch, но я ищу реализацию в numpy. Я использовал where , но, похоже, это не работает, также возможно ли объединить поиск макса и индексов за один раз, чтобы быть более эффективным? Спасибо за любую помощь!

У меня есть решение, аналогичное решению Andras, основанному на np.argmax и np.arange. Вместо «индексирования индекса» я предлагаю добавить кусочное смещение к результату np.argmax:

 import numpy as np
a = np.array([[[7, 9],
               [19, 18]],
              [[24, 5],
               [18, 11]]])
off = np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])
 

 >>> off
array([[0, 2],
       [4, 6]])
 

Это приводит к:

 >>> a.argmax(-1)   off
array([[1, 2],
       [4, 6]])
 

Или в качестве однострочного:

 >>> a.argmax(-1)   np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])
array([[1, 2],
       [4, 6]])
 

Единственное решение, которое я мог бы придумать прямо сейчас, — это создание 2d (или 3d, см. Ниже) range , которое индексирует ваш плоский массив, и индексирование в него с максимальными индексами, которые определяют b (т. Е. a.argmax(-1) ):

 import numpy as np

a = np.array([[[ 7,  9],
               [19, 18]],
              [[24,  5],
               [18, 11]]])
multi_inds = a.argmax(-1)
b_shape = a.shape[:-1]
b_size = np.prod(b_shape)
flat_inds = np.arange(a.size).reshape(b_size, -1)
flat_max_inds = flat_inds[range(b_size), multi_inds.ravel()]
max_inds = flat_max_inds.reshape(b_shape)
 

Я разделил шаги некоторыми значимыми именами переменных, которые, надеюсь, должны объяснить, что происходит.

multi_inds сообщает вам, какой «столбец» выбрать в каждой «строке», a чтобы получить максимум:

 >>> multi_inds
array([[1, 0],
       [0, 0]])
 

flat_inds представляет собой список индексов, из которых в каждой строке должно быть выбрано одно значение:

 >>> flat_inds
array([[0, 1],
       [2, 3],
       [4, 5],
       [6, 7]])
 

Это индексируется точно в соответствии с максимальными индексами в каждой строке. flat_max_inds являются ли значения, которые вы ищете, но в плоском массиве:

 >>> flat_max_inds
array([1, 2, 4, 6])
 

Поэтому нам нужно изменить это обратно, чтобы соответствовать b.shape :

 >>> max_inds
array([[1, 2],
       [4, 6]])
 

Немного более неясным, но и более элегантным решением является использование массива 3d-индексов и использование в нем широковещательной индексации:

 import numpy as np

a = np.array([[[ 7,  9],
               [19, 18]],
              [[24,  5],
               [18, 11]]])
multi_inds = a.argmax(-1)
i, j = np.indices(a.shape[:-1])
max_inds = np.arange(a.size).reshape(a.shape)[i, j, multi_inds]
 

Это делает то же самое без промежуточного выравнивания в 2d.

Последняя часть также заключается в том , как вы можете получить b multi_inds , т. Е. Без необходимости вызывать *max функцию во второй раз:

 b = a[i, j, multi_inds]
 

Это длинный однострочный

 new = np.array([np.where(a.reshape(-1)==x)[0][0] for x in a.max(-1).reshape(-1)]).reshape(2,2)
 

 print(new)
array([[1, 2],
       [4, 3]])
 

Однако число = 18 повторяется дважды; Так какой индекс является целевым.