返回非平坦指数的numpy数组的Argmax

时间:2012-02-28 13:12:05

标签: python multidimensional-array numpy

我正在尝试获取Numpy数组中最大元素的索引。 这可以使用numpy.argmax完成。我的问题是,我希望找到整个数组中最大的元素并得到它的索引。

numpy.argmax可以沿着一个轴应用,这不是我想要的,也可以应用在扁平数组上,这是我想要的。

我的问题是,当我需要多维索引时,使用numpy.argmaxaxis=None会返回平面索引。

我可以使用divmod来获得非平坦的索引,但这感觉很难看。有没有更好的方法呢?

3 个答案:

答案 0 :(得分:120)

您可以对numpy.argmax()

的结果使用numpy.unravel_index()
>>> a = numpy.random.random((10, 10))
>>> numpy.unravel_index(a.argmax(), a.shape)
(6, 7)
>>> a[6, 7] == a.max()
True

答案 1 :(得分:16)

np.where(a==a.max())

返回最大元素的坐标,但必须解析数组两次。

>>> a = np.array(((3,4,5),(0,1,2)))
>>> np.where(a==a.max())
(array([0]), array([2]))

argmax相比,这会返回所有元素的坐标等于最大值。 argmax只返回其中一个(np.ones(5).argmax()返回0)。

答案 2 :(得分:1)

要获取所有出现的最大值的非平坦索引,可以使用argwhere代替where来稍微修改eumiro的answer

np.argwhere(a==a.max())

>>> a = np.array([[1,2,4],[4,3,4]])
>>> np.argwhere(a==a.max())
array([[0, 2],
       [1, 0],
       [1, 2]])