numpy:如何从argmax结果中获取最大值 [英] numpy: how to get a max from an argmax result

问题描述

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

I have a numpy array of arbitrary shape, e.g.:

a = array([[[ 1,  2],
            [ 3,  4],
            [ 8,  6]],

          [[ 7,  8],
           [ 9,  8],
           [ 3, 12]]])
a.shape = (2, 3, 2)

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

and a result of argmax over the last axis:

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

我想得到最大:

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

但无需重新计算所有内容.我尝试过:

But without recalculating everything. I've tried:

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) 

该怎么做?二维的类似问题: numpy二维二维max/argmax

How to do it? Similar question for 2-d: numpy 2d array max/argmax

推荐答案

您可以使用

You can use advanced indexing -

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情况，如 np.ix_ ，就像这样-

For a generic ndarray case of any number of dimensions, as stated in the comments by @hpaulj, we could use np.ix_, like so -

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:如何从argmax结果中获取最大值的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！