沿某个轴的多维数组的 argmax

Question

我有一个看起来像这样的多维数组：

my_array = np.arange(2)[:,None,None] *np.arange(4)[:, None]*np.arange(8)

我正在寻找二维 argmax 的多维等价物

特别是，我正在寻找沿轴 = 2 的最大值的 argmax。我先尝试整形，但整形会完全破坏整个数组的原始索引信息，所以它可能不起作用。我不知道该怎么做，需要你们的帮助。提前谢谢你

编辑：期望输出是：

[(0,0,0),(1,3,1),(1,3,2),(1,3,3),(1,3,4),(1,3,5),(1,3,6),(1,3,7)]

这正是沿轴 = 2 的最大值索引的数组

Answer 1

为了在 3D ndarray 的最后一个轴上找到这样的 argmax 索引，我们可以使用沿着这些线的东西 -

In [66]: idx = my_array.reshape(-1,my_array.shape[-1]).argmax(0)

In [67]: r,c = np.unravel_index(idx,my_array.shape[:-1])

In [68]: l = np.arange(len(idx))

In [69]: np.c_[r,c,l]
Out[69]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])

将其扩展到通用 ndarray -

In [99]: R = np.unravel_index(idx,my_array.shape[:-1])

In [104]: np.hstack((np.c_[R],l[:,None]))
Out[104]: 
array([[0, 0, 0],
       [1, 3, 1],
       [1, 3, 2],
       [1, 3, 3],
       [1, 3, 4],
       [1, 3, 5],
       [1, 3, 6],
       [1, 3, 7]])