获取轴上numpy.argmax元素的索引

时间:2013-11-21 18:16:53

标签: python arrays numpy indexing argmax

我有N维矩阵,其中包含具有N个参数的函数的值。每个参数都有一个离散数量的值。我需要最大化除了一个参数之外的所有参数的函数,从而产生大小等于非最大化参数的值的一维向量。我还需要保存其他参数所采用的值。

为此,我想在不同的轴上迭代地应用numpy.max以减少矩阵的维数以找到我需要的东西。最终的向量将取决于我遗漏的参数。

然而,我无法找到最终元素的原始索引(其中包含有关其他参数所采用值的信息)。我虽然以与numpy.argmax相同的方式使用numpy.max,但我无法获得原始索引。

我正在尝试的一个例子是:

x = [[[1,2],[0,1]],[[3,4],[6,7]]]
args = np.argmax(x, 0)

返回

[[1 1]
 [1 1]]

这意味着argmax正在选择原始矩阵中的元素(2,1,4,7)。但如何获得他们的指数?我尝试unravel_index,直接使用args作为矩阵x的索引,从numpy到索引的一堆函数没有成功。

使用numpy.where不是解决方案,因为输入矩阵内部可能具有相同的值,因此我无法识别不同的原始值。

1 个答案:

答案 0 :(得分:9)

x.argmax(0)为第1轴提供最大值的索引。使用np.indices生成另一个轴的索引。

x = np.array([[[1,2],[0,1]],[[3,4],[6,7]]])
x.argmax(0)
    array([[1, 1],
           [1, 1]])
a1, a2 = np.indices((2,2))
(x.argmax(0),a1,a2)
    (array([[1, 1],
            [1, 1]]),
     array([[0, 0],
            [1, 1]]),
     array([[0, 1],
            [0, 1]]))


x[x.argmax(0),a1,a2]
    array([[3, 4],
           [6, 7]])

x[a1,x.argmax(1),a2] 
    array([[1, 2],
           [6, 7]])

x[a1,a2,x.argmax(2)] 
    array([[2, 1],
           [4, 7]])

如果x包含其他维度,请相应地生成a1a2

官方文档没有说明如何使用argmax,但早期的SO线程已经讨论过它。我从Using numpy.argmax() on multidimensional arrays

得到了这个大致的想法