如何在numpy数组中获取N个最大值的索引？

2018-06-18 01:37:38

Numpy提出了一种通过np.argmax数组最大值索引的np.argmax 。

我想要一个类似的东西，但返回N个最大值的索引。

例如，如果我有一个数组[1, 3, 2, 4, 5] ，它的function(array, n=3)将返回[4, 3, 1] 。

谢谢：）

我能够想到的最简单的是：

In [1]: import numpy as np

In [2]: arr = np.array([1, 3, 2, 4, 5])

In [3]: arr.argsort()[-3:][::-1]
Out[3]: array([4, 3, 1])

这涉及到一个完整的数组。我不知道numpy提供了一种内置的方式来进行部分排序; 到目前为止，我还没有找到一个。

如果这个解决方案变得太慢了（特别是对于小n ），那么可能值得用Cython编码。

较新的NumPy版本（1.8及以上版本）有一个称为argpartition的功能。要获得四个最大元素的指标，请执行

>>> a
array([9, 4, 4, 3, 3, 9, 0, 4, 6, 0])
>>> ind = np.argpartition(a, -4)[-4:]
>>> ind
array([1, 5, 8, 0])
>>> a[ind]
array([4, 9, 6, 9])

与argsort不同的argsort ，这个函数在最坏的情况下以线性时间运行，但是返回的索引没有被排序，从评估a[ind]的结果可以看出。如果您也需要，请在之后进行分类：

>>> ind[np.argsort(a[ind])]
array([1, 8, 5, 0])

以这种方式获得排序顺序的top-k元素需要O（n + k log k）时间。

编辑：修改包括Ashwini Chaudhary的改进。

>>> import heapq
>>> import numpy
>>> a = numpy.array([1, 3, 2, 4, 5])
>>> heapq.nlargest(3, range(len(a)), a.take)
[4, 3, 1]

对于常规Python列表：

>>> a = [1, 3, 2, 4, 5]
>>> heapq.nlargest(3, range(len(a)), a.__getitem__)
[4, 3, 1]

如果您使用Python 2，请使用xrange而不是range 。

资料来源：http://docs.python.org/3/library/heapq.html

链接地址: http://www.djcxy.com/p/50985.html