python - 有效地查找 Python Array/List 中 N 个最大元素的索引

Question

如果这是一个重复的问题，我很抱歉，我查找了此信息但仍然找不到。

是否可以通过使用N个最大元素的索引以降序非常有效地排列一个numpy数组(或python列表)？

比如数组:

a = array([4, 1, 0, 8, 5, 2])

按降序排列的最大元素的索引将给出(考虑到 N = 6，包括所有元素):

8 --> 3

5 --> 4

4 --> 0

2 --> 5

1 --> 1

0 --> 2

result = [3, 4, 0, 5, 1, 2]

我知道如何使用一种有点愚蠢的方法来实现它(比如对数组进行排序并搜索 N 个数字中的每一个作为它们的索引)，但我想知道是否有任何有效的库，如瓶颈或 heapq 或者 pythonic使这个非常快的方法。我必须将它应用到多个数组中，每个数组包含 300k 个元素，这就是性能成为问题的原因。

提前致谢!

更新

我阅读了答案并决定使用 300k 的随机整数对它们进行计时，结果如下:

解决方案 1: sorted(range(len(a)), key=lambda i:a[i]) 时间: 230毫秒

解决方案2: heapq.nlargest(len(a), zip(a, itertools.count())) 时间: 396毫秒

解决方案 3: heapq.nlargest(len(a), enumerate(a), key=operator.itemgetter(1)) 时间: 864 毫秒

解决方案4: def f(a,N): return np.argsort(a)[::-1][:N] (N = len(a)) 时间:104 毫秒

非常感谢您提供快速且非常好的答案!

Answer 1

你看过numpy内置的argsort方法吗？:

http://docs.scipy.org/doc/numpy/reference/generated/numpy.argsort.html

我可以在我的机器上使用该方法在大约 29 毫秒内对包含 300,000 个随机 float 的数组进行排序。

def f(a,N):
    return np.argsort(a)[::-1][:N]