Получить индексы элемента одного массива, используя индексы в другом массиве

#python #numpy #indexing

Вопрос:

Предположим, у меня есть массив 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 , но, похоже, это не работает, также возможно ли объединить поиск макса и индексов за один раз, чтобы быть более эффективным? Спасибо за любую помощь!

Ответ №1:

У меня есть решение, аналогичное решению 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]])
 

Комментарии:

1. вау, спасибо @Per Joachims, это тоже работает одинаково хорошо! Поскольку вы «Новый участник», я думаю, что было бы лучше принять ваш ответ, лол. Не могли бы вы, пожалуйста, немного объяснить смещение? Я действительно не понимаю этого… Спасибо!!

2. Я думаю, что понимаю это: таким образом, в основном off дается индекс первого элемента в каждой строке, затем плюс argmax даст индекс максимального элемента в этой строке. Это просто, еще раз спасибо

Ответ №2:

Единственное решение, которое я мог бы придумать прямо сейчас, — это создание 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]
 

Комментарии:

1. вау, спасибо, Андрас. это выглядит довольно сложно. Знаете ли вы какие-либо функции numpy, которые возвращают как максимальное значение, так и индексы? argmax в основном снова находит макса, верно?

2. @Sam-gege мой обновленный ответ менее сложен, но его также труднее понять. Вы можете просто позвонить argmax и запросить значения, так же, как я сделал с диапазоном в обновленном ответе. b = a[i, j, multi_inds] должно хватить. (Я тоже обновил свой ответ этим.)

3. классно, спасибо за помощь! Мне нужно попробовать это

4. спасибо за ваш обновленный ответ. Я думаю, что мне нужно позвонить max хотя бы один раз в maxpool, но я также хочу получить его индексы, чтобы во время обратного распространения я мог вернуть градиент в положение max. интересно, поскольку массив numpy является непрерывным, функции max() не должно быть очень сложно также возвращать индекс в соответствии с выравниванием входных данных, но, к сожалению, это не так . Еще раз спасибо за вашу помощь!!

5. да, вы правы, я забыл об этом, так что мне нужно позвонить в argmax только один раз. после того, как у меня будут эти индексы, я могу просто позвонить a.reshape(-1)[np.array([[1, 2],[4, 6]])] , чтобы получить эти максимальные значения. в принципе, мне не нужно звонить максу, а затем argmax

Ответ №3:

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

 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 повторяется дважды; Так какой индекс является целевым.

Комментарии:

1. Спасибо за ваш ответ, дамир. Однако это неверно, так как должен быть возвращен индекс вторых 18, так как эти 18 являются максимальными в своей строке (18,11), а не первыми 18 в строке (19,18)