python - 在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 是以下矩阵:

[[[ 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))

我可以得到矩阵结果:

[[ 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]]

，但是我怎样才能获得每个最大值的整个索引，比如 [0,0,1] 之类的，因为我需要索引来绘制 3D 图形。

Answer 1

如果我理解正确，您希望获得沿最后两个轴的最大值的索引。可以这样完成

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

产量

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]])