python - 在 numpy 的 3d 数组中。如何提取三维中最大元素的索引？

Question

示例输入形状为 (2,2,2) 的 3D 数组:

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

我的 3d 数组的形状为 (N, N, N)，在上面的示例中 N = 2。

我需要获取所有索引，使得第三维的索引属于第三维的最大元素，以上 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 的方法获取沿第一轴和第二轴的所有索引，然后使用 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()))

sample 运行-

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