Python Numpy:argmax在不带循环的多个轴上

我有一个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)

在这方面，我们将：

A: 
[[[ 0  1]
  [ 2  3]]

 [[ 4  5]
  [ 6  7]]

 [[ 8  9]
  [10 11]]

 [[ 0 -1]
  [-2 -3]]]

maxpos:
[[ 1.  1.]
 [ 1.  1.]
 [ 1.  1.]
 [ 0.  0.]]

如果有多个最大化器，我不介意选择哪个

---

我曾尝试使用

np.apply\u over\u axes

，但我没有成功使其返回我想要的结果。

您可以这样做-

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

def argmax_lastNaxes(A, N):
    s = A.shape
    new_shp = s[:-N] + (np.prod(s[-N:]),)
    max_idx = A.reshape(new_shp).argmax(-1)
    return np.unravel_index(max_idx, s[-N:])

样本运行-

In [214]: # Input array
     ...: A = np.random.rand(5,4,3,7,8)

In [215]: # Setup output array and use original loopy code
     ...: maxpos=np.empty(shape=(5,4)) # 4 because ndims in A is 5
     ...: for n in range(0, 5):
     ...:     maxpos[n,:]=np.unravel_index(np.argmax(A[n,:,:,:,:]), A[n,:,:,:,:].shape)
     ...:     

In [216]: # Proposed approach
     ...: max_idx = A.reshape(A.shape[0],-1).argmax(1)
     ...: maxpos_vect = np.column_stack(np.unravel_index(max_idx, A[0,:,:].shape))
     ...: 

In [219]: # Verify results
     ...: np.array_equal(maxpos.astype(int),maxpos_vect)
Out[219]: True

n-dim阵列的推广 我们可以将n-dim数组的解推广到最后一个

n

轴的

argmax

，并结合类似的内容-

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

def argmax_lastNaxes(A, N):
    s = A.shape
    new_shp = s[:-N] + (np.prod(s[-N:]),)
    max_idx = A.reshape(new_shp).argmax(-1)
    return np.unravel_index(max_idx, s[-N:])

---

结果将是索引数组的元组。如果您需要作为数组的最终输出，我们可以使用

np.stack

或

np.concatenate

您可以使用列表

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

它的可读性更高，但速度不会比显式循环好多少

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

如果是这种情况（即，您有一千万个2x2矩阵），则翻转速度更快

# true if 0,0 is not smaller than others
m00 = ((data[:,0,0] >= data[:,1,0]) &
       (data[:,0,0] >= data[:,0,1]) &
       (data[:,0,0] >= data[:,1,1]))

# true if 0,1 is not smaller than others
m01 = ((data[:,0,1] >= data[:,1,0]) &
       (data[:,0,1] >= data[:,0,0]) &
       (data[:,0,1] >= data[:,1,1]))

# true if 1,0 is not smaller than others
m10 = ((data[:,1,0] >= data[:,0,0]) &
       (data[:,1,0] >= data[:,0,1]) &
       (data[:,1,0] >= data[:,1,1]))

# true if 1,1 is not smaller than others
m11 = ((data[:,1,1] >= data[:,1,0]) &
       (data[:,1,1] >= data[:,0,1]) &
       (data[:,1,1] >= data[:,0,0]))

# choose which is max on equality
m01 &= ~m00
m10 &= ~(m00|m01)
m11 &= ~(m00|m01|m10)

# compute result
result = np.zeros((len(data), 2), np.int32)
result[:,1] |= m01|m11
result[:,0] |= m10|m11

---

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

这个答案已链接到

maxpos

是OP要求的，但它如何有用？@hpaulj不确定“有用性”是什么意思。Maxpos值是最后两个轴的argmax值，用于三维阵列输入，索引引用回最后两个轴的形状。这回答了你的问题吗？或者您想知道这是如何应用于n-dim阵列的？