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

Question

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

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

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

并且在最后一个轴上进行argmax操作的结果：

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

我想取得最大值：

np.max(a, axis=-1) = array([[ 2,  4,  8],
                            [ 8,  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) 

如何操作？2-d类似的问题：numpy 2d数组的最大值/最大值索引

- sygi

---

重塑后的_x是什么？ - dnalow

抱歉，应该是 a。现在正在更正。 - sygi

Answer 1

您可以使用高级索引 -

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情况，如在@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]

- Divakar

---

好的，但如何使用任意形状的数组而不仅仅是3维呢？ - sygi

使用np.ix_生成第一个维度，并使用元组连接idx。 - hpaulj

@hpaulj 谢谢，那就可以了！已经更新了帖子。 - Divakar

谢谢。有错别字吗？似乎dims没有被使用。 - sygi

@sygi 抱歉，我试了几个东西，最后发现那不是必需的。已经删除了那一行。 - Divakar

FutureWarning 表示我们想要元组 a[tuple([*np.ix_(*[np.arange(i) for i in a.shape[:-1]]), idxs])]。 - Franz Knülle

Answer 2

对于任意形状的数组，以下代码应该可以工作:)

a = np.arange(5 * 4 * 3).reshape((5,4,3))

# for last axis
argmax = a.argmax(axis=-1)
a[tuple(np.indices(a.shape[:-1])) + (argmax,)]

# for other axis (eg. axis=1)
argmax = a.argmax(axis=1)
idx = list(np.indices(a.shape[:1]+a.shape[2:]))
idx[1:1] = [argmax]
a[tuple(idx)]

或者

a = np.arange(5 * 4 * 3).reshape((5,4,3))

argmax = a.argmax(axis=0)
np.choose(argmax, np.moveaxis(a, 0, 0))

argmax = a.argmax(axis=1)
np.choose(argmax, np.moveaxis(a, 1, 0))

argmax = a.argmax(axis=2)
np.choose(argmax, np.moveaxis(a, 2, 0))

argmax = a.argmax(axis=-1)
np.choose(argmax, np.moveaxis(a, -1, 0))

- lianwei82cn

Answer 3

对于任意形状的ndarray，您可以将argmax索引展平，然后恢复正确的形状，如下所示：

idx = np.argmax(a, axis=-1)
flat_idx = np.arange(a.size, step=a.shape[-1]) + idx.ravel()
maximum = a.ravel()[flat_idx].reshape(*a.shape[:-1])

- DavidG