在python中的3维矩阵中查找每一行的最大值和索引

Question

import numpy as np

a=np.zeros(shape=(3,4,5))
print(type(a))


a[0][1]=1,2,3,4,5

a[1][3]=6,7,8,9,12

a[2][1]=12,34,54,21,89

a is the following matrix:

[[[ 0.  0.  0.  0.  0.]
[ 1.  2.  3.  4.  5.]
[ 0.  0.  0.  0.  0.]
[ 0.  0.  0.  0.  0.]]

[[ 0.  0.  0.  0.  0.]
[ 0.  0.  0.  0.  0.]
[ 0.  0.  0.  0.  0.]
[ 6.  7.  8.  9. 12.]]

[[ 0.  0.  0.  0.  0.]
[12. 34. 54. 21. 89.]
[ 0.  0.  0.  0.  0.]
[ 0.  0.  0.  0.  0.]]]

.

print(a.max(axis=0))
print(a.argmax(axis=0))

I can get the matrix results:

[[ 0.  0.  0.  0.  0.]
[12. 34. 54. 21. 89.]
[ 0.  0.  0.  0.  0.]
[ 6.  7.  8.  9. 12.]]
[[0 0 0 0 0]
[2 2 2 2 2]
[0 0 0 0 0]
[1 1 1 1 1]]

, but how can I get the whole index of each maximum like [0,0,1] or something, as I need the index to draw a 3D figure.

Answer 1

If I understood you correctly, you'd like to get the index of the maximum along the last two axes. This can be done like this

np.array([[i, j, np.argmax(a[i, j])] for i in range(a.shape[0]) for j in range(a.shape[1])])

yielding

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