Question

示例输入形状（2,2,2）的3D数组：

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

我的3d数组的形状为（N，N，N），在上例中N = 2。

我需要得到所有索引，使得第三维的索引属于第三维的max元素，输出为上面的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()))

示例运行 -

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

在numpy的3d数组中。如何在第三维中提取max元素的索引？

1 个答案: