numpy：如何从argmax结果中获取最大值

问题描述

您可以使用advancedindexing-

In [17]: a
Out[17]: 
array([[[ 1,  2],
        [ 3,  4],
        [ 8,  6]],

       [[ 7,  8],
        [ 9,  8],
        [ 3, 12]]])

In [18]: idx = a.argmax(axis=-1)

In [19]: m,n = a.shape[:2]

In [20]: a[np.arange(m)[:,None],np.arange(n),idx]
Out[20]: 
array([[ 2,  4,  8],
       [ 8,  9, 12]])

对于任何数量的维度的通用ndarray情况，如中所述comments by@hpaulj，我们可以使用np.ix_，就像这样-

shp = np.array(a.shape)
dim_idx = list(np.ix_(*[np.arange(i) for i in shp[:-1]]))
dim_idx.append(idx)
out = a[dim_idx]

解决方法

我有一个任意形状的numpy数组，例如：

a = array([[[ 1,2],[ 3,4],[ 8,6]],[[ 7,8],[ 9,12]]])
a.shape = (2,3,2)

以及最后一个轴上argmax的结果：

np.argmax(a,axis=-1) = array([[1,1,0],[1,1]])

我想得到最大：

np.max(a,axis=-1) = array([[ 2,4,9,12]])

但无需重新计算所有内容。我试过了：

a[np.arange(len(a)),np.argmax(a,axis=-1)]

但是得到了：

IndexError: shape mismatch: indexing arrays could not be broadcast together with shapes (2,) (2,3)