使用另一个数组中的索引获取一个数组元素的索引

Question

假设我有一个形状为 (2, 2, 2) 的数组 a：

a = np.array([[[7, 9],
               [19, 18]],
              [[24, 5],
               [18, 11]]])

还有一个数组b，它是a 的最大值：b=a.max(-1)（逐行）：

b = np.array([[9, 19],
              [24, 18]])

我想使用扁平化a 中的索引获取b 中元素的索引，即a.reshape(-1)：

array([ 7,  9, 19, 18, 24,  5, 18, 11])

结果应该是一个与b 形状相同的数组，其索引为b 在扁平a 中：

array([[1, 2],
       [4, 6]])

基本上这是在pytorch中return_indices = True时maxpool2d的结果，但我正在寻找numpy中的实现。我使用过where，但它似乎不起作用，是否可以一次将查找最大值和索引结合起来，以提高效率？感谢您的帮助！

Answer 1

我现在能想到的唯一解决方案是生成一个 2d（或 3d，见下文）range 来索引您的平面数组，并使用定义 b 的最大索引（即 a.argmax(-1) ):

import numpy as np

a = np.array([[[ 7,  9],
               [19, 18]],
              [[24,  5],
               [18, 11]]])
multi_inds = a.argmax(-1)
b_shape = a.shape[:-1]
b_size = np.prod(b_shape)
flat_inds = np.arange(a.size).reshape(b_size, -1)
flat_max_inds = flat_inds[range(b_size), multi_inds.ravel()]
max_inds = flat_max_inds.reshape(b_shape)

我用一些有意义的变量名分隔了这些步骤，希望能解释发生了什么。

multi_inds 告诉您在a 的每一“行”中选择哪个“列”以获得最大值：

>>> multi_inds
array([[1, 0],
       [0, 0]])

flat_inds 是一个索引列表，每行从中选择一个值：

>>> flat_inds
array([[0, 1],
       [2, 3],
       [4, 5],
       [6, 7]])

这完全根据每行中的最大索引进行索引。 flat_max_inds 是您要查找的值，但在一个平面数组中：

>>> flat_max_inds
array([1, 2, 4, 6])

所以我们需要重新整形以匹配b.shape：

>>> max_inds
array([[1, 2],
       [4, 6]])

---

一个稍微晦涩但也更优雅的解决方案是使用 3d 索引数组并使用广播索引：

import numpy as np

a = np.array([[[ 7,  9],
               [19, 18]],
              [[24,  5],
               [18, 11]]])
multi_inds = a.argmax(-1)
i, j = np.indices(a.shape[:-1])
max_inds = np.arange(a.size).reshape(a.shape)[i, j, multi_inds]

这做同样的事情，没有中间展平成 2d。

最后一部分也是如何从multi_inds 获取b，即无需再次调用*max 函数：

b = a[i, j, multi_inds]

Answer 2

这是一条长线

new = np.array([np.where(a.reshape(-1)==x)[0][0] for x in a.max(-1).reshape(-1)]).reshape(2,2)

print(new)
array([[1, 2],
       [4, 3]])

但是 number = 18 重复了两次；那么目标是哪个索引。

Answer 3

我有一个类似于 Andras 基于 np.argmax 和 np.arange 的解决方案。我建议在 np.argmax 的结果中添加 分段偏移，而不是“索引索引”：

import numpy as np
a = np.array([[[7, 9],
               [19, 18]],
              [[24, 5],
               [18, 11]]])
off = np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])

>>> off
array([[0, 2],
       [4, 6]])

这会导致：

>>> a.argmax(-1) + off
array([[1, 2],
       [4, 6]])

或者作为单行：

>>> a.argmax(-1) + np.arange(0, a.size, a.shape[2]).reshape(a.shape[0], a.shape[1])
array([[1, 2],
       [4, 6]])