В 3d массив в numpy. Как я могу извлечь индексы максимального элемента в третьем измерении?

Question

Пример входных данных 3D-массив формы (2,2,2):

[[[ 1, 2], 
    [ 4, 3]], 
[[ 5, 6], 
    [ 8, 7]]] 

Мой 3d массив имеет форму (N, N, N), в приведенном выше примере N = 2.

I нужно, чтобы все показатели, такие, что индекс третьего измерения относится к максимальному элементу в третьем измерении, выход для 3D выше массива:

[[0, 0, 1], # for element 2 
[0, 1, 0], # for element 4 
[1, 0, 1], # for element 6 
[1, 1, 0]] # for element 8 

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

Answer 1

Вот подход, использующий np.meshgrid, чтобы получить все индексы вдоль первой и второй осей, а затем складывая их Alongwith на максимальных показателей от третьей оси с помощью np.column_stack -

d = a.argmax(-1) 
m,n = a.shape[:2] 
c,r = np.mgrid[:m,:n] 
out = np.column_stack((c.ravel(),r.ravel(),d.ravel())) 

Пример запуска -

In [96]: a 
Out[96]: 
array([[[38, 49, 15, 61, 29], 
     [31, 88, 45, 88, 20], 
     [17, 97, 58, 61, 14], 
     [43, 77, 56, 92, 89]], 

     [[48, 91, 49, 35, 58], 
     [53, 34, 58, 92, 52], 
     [20, 35, 70, 41, 81], 
     [60, 42, 85, 82, 41]], 

     [[45, 41, 32, 41, 25], 
     [59, 32, 90, 18, 47], 
     [24, 93, 29, 89, 12], 
     [80, 27, 12, 51, 33]]]) 

In [97]: out 
Out[97]: 
array([[0, 0, 3], 
     [0, 1, 1], 
     [0, 2, 1], 
     [0, 3, 3], 
     [1, 0, 1], 
     [1, 1, 3], 
     [1, 2, 4], 
     [1, 3, 2], 
     [2, 0, 0], 
     [2, 1, 2], 
     [2, 2, 1], 
     [2, 3, 0]]) 

---

В качестве альтернативы, поскольку эти индексы в основном являются повторениями, мы можем использовать np.repeat и np.tile, чтобы получить эти индексы массивов, а затем использовать np.column_stack, как и прежде, как так -

d0 = np.arange(m).repeat(n) 
d1 = np.tile(np.arange(n),m) 
out = np.column_stack((d0,d1,d.ravel()))