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

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

您可以使用

advancedindexing

-

In [17]: aOut[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]