python - 在定义的轴上使用最大值索引

Question

我想在矩阵“A”的每一行的最大值的相同索引处提取数组“B”的值。为了找到索引，我使用了 numpy 函数“numpy.argmax”，如下所示：

>>> A=numpy.asarray([[0,1,6],[3,2,4]]);A
array([[0, 1, 6],
       [3, 2, 4]])
>>> argA=numpy.argmax(A,axis=1);argA
array([2, 2])

问题是我不知道如何使用“argA”来提取数组“B”中的值

Answer 1

中的每个条目argA对应于相应行内最大值的索引位置。这些行不是显式的（由于使用axis=1），而是对应于每个条目的索引。因此，您需要添加它们以获取您所追求的元素。

>>> A[[0,1], argA]
array([6, 4])

所以：

>>> B
array([[ 9,  8,  2],
       [ 3,  4,  5]])
>>> B[[0,1], argA] = 84,89
>>> B
array([[ 9,  8, 84],
       [ 3,  4, 89]])

概括使用：

>>> B[np.arange(A.shape[0]),argA]

Answer 2

您可以直接使用数组作为索引 - 这实际上很简单：

B[:, arga]

例如：

>>> A[:,argA]
array([[6, 6],
       [4, 4]])

Answer 3

这可能看起来有点低效，但您可以使用：

A.take(argA, axis=1).diagonal()

而不是A.take你可以使用A[:, argA]但take更明确的轴。

您可以使用ravel_multi_index转换为平面索引：

A.take(np.ravel_multi_index((np.arange(len(argA)), argA), A.shape))