argmax для многомерного массива вдоль некоторой оси - PullRequest
Новая
       10

argmax для многомерного массива вдоль некоторой оси

1 голос
/ 05 мая 2019

У меня есть многомерный массив, который выглядит так: 

my_array = np.arange(2)[:,None,None] *np.arange(4)[:, None]*np.arange(8)

Я ищу многомерный эквивалент 2-D argmax

В частности, я ищу argmax максимумов вдоль оси = 2. Сначала я попытался изменить форму, но изменение формы полностью уничтожит информацию об исходных индексах всего массива, поэтому, вероятно, она не будет работать. Я понятия не имею, как это сделать, и вам нужна помощь от вас, ребята. Заранее спасибо

РЕДАКТИРОВАТЬ: Желание выхода: 

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

Это точно массив индексов максимумов по оси = 2

1 Ответ

1 голос
/ 05 мая 2019

Для нахождения таких индексов argmax вдоль последней оси трехмерного массива мы можем использовать что-то по этим линиям - 

In [66]: idx = my_array.reshape(-1,my_array.shape[-1]).argmax(0)

In [67]: r,c = np.unravel_index(idx,my_array.shape[:-1])

In [68]: l = np.arange(len(idx))

In [69]: np.c_[r,c,l]
Out[69]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])

Чтобы распространить это на общий массив ndarray - 

In [99]: R = np.unravel_index(idx,my_array.shape[:-1])

In [104]: np.hstack((np.c_[R],l[:,None]))
Out[104]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])
...