Numpy:argmax在多个轴上无循环

Question

我有一个N维数组（名为a）。对于A的第一个轴的每一行，我想获得沿着A的其他轴的最大值的坐标。然后我将返回一个二维数组，其中包含A的第一个轴的每一行的最大值的坐标

我已经用一个循环解决了我的问题，但我想知道是否有一种更有效的方法来实现这一点。我当前的解决方案（例如数组A）如下：

import numpy as np

A=np.reshape(np.concatenate((np.arange(0,12),np.arange(0,-4,-1))),(4,2,2))
maxpos=np.empty(shape=(4,2))
for n in range(0, 4):
    maxpos[n,:]=np.unravel_index(np.argmax(A[n,:,:]), A[n,:,:].shape)

在这里，我们可以：

^{pr2}$

如果有多个最大化者，我不介意选择哪一个。在

我尝试过使用np.apply_over_axes，但是我没能让它返回我想要的结果。在

Answer 1

你可以这样做-

# Reshape input array to a 2D array with rows being kept as with original array.
# Then, get idnices of max values along the columns.
max_idx = A.reshape(A.shape[0],-1).argmax(1)

# Get unravel indices corresponding to original shape of A
maxpos_vect = np.column_stack(np.unravel_index(max_idx, A[0,:,:].shape))

样本运行-

^{pr2}$

Answer 2

你可以使用列表理解

result = [np.unravel_index(np.argmax(r), r.shape) for r in a]

它在我看来更具可读性，但速度不会比显式循环好多少。在

只有当第一个维度实际上非常大时，主外循环在Python中才有意义。在

如果是这种情况（例如，你有一千万个2x2矩阵），那么翻转速度更快。。。在

^{pr2}$

在我的机器上，上面的代码大约快50倍（对于一百万个2x2矩阵）。在