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

【问题标题】：Finding the Index of N biggest elements in Python Array / List Efficiently有效地查找 Python Array / List 中 N 个最大元素的索引
【发布时间】：2012-09-29 01:17:03
【问题描述】：

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

是否可以通过使用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 毫秒

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

【问题讨论】：

How to get the N maximum values in a numpy array? 的可能重复项
如果你跟踪重复的踪迹，this answer 会弹出，这似乎很有希望——尽管该帖子是由开发人员发布的，但答案并未透露这一事实......
在您的测试中， N 的值是多少？如上所述，使用 heapq 是有效的，因为 N 与 len(a) 相比相当小。
你如何为N < len(a)修改这些？
我同意@lizzie。您能否在实验中提供N 和len(a) 的值？如果N 比len(a) 小很多，我认为heapq.nlargest 应该比np.argsort 更有效。

标签： python performance numpy

【解决方案1】：

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

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

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

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

【讨论】：

这效果非常好！在我的机器上它需要 104 毫秒（现在很忙），稍后我会再试一次，但到目前为止，这是最快的解决方案。 Tnx！
@joshadel 这个函数首先是 argsort，然后返回前 N 个值。是否有一个 Numpy/Scipy 函数等效于 Python heapq.nlargest(N, a) 但无需对整个数组进行 argsort 即可查找前 N 个索引值？
@dbv 可能类似于berkeleyanalytics.com/bottleneck/…，但这似乎只适用于最小值。

【解决方案2】：

L = [4, 1, 0, 8, 5, 2]
sorted(range(len(L)), key=lambda i:L[i])

【讨论】：

比我聪明。 +1 先生。
key = L.__getitem__ 是另一种选择（在某些情况下可能会更快一些）。
@GarethRees：你是对的！我没有想到这一点。 lambdas /are/ 慢
我尝试了一个简单的测试，并没有发现太大的差异，所以使用lambda 不会错。
我尝试使用 getitem 使其工作，但由于我是 python 的新手，无法使其正常工作，但使用 lambda 的解决方案在这里运行良好，感谢您的帮助！

【解决方案3】：

您可以使用heapq 轻松做到这一点：

>>> heapq.nlargest(3, zip(a, itertools.count()))
[(8, 3), (5, 4), (4, 5)]

元组是通过对第一个值排序，然后是第二个，等等... ，但我们可以很容易地把它们扔掉）。

我使用zip() 和itertools.count() 作为枚举给了我们错误的顺序，所以它们将按索引排序，而不是按值排序。或者，您也可以使用((value, index) for index, value in enumerate(a))，但我觉得不太清楚。

另一种选择是给一个密钥，做heapq.nlargest(3, enumerate(a), key=operator.itemgetter(1))。

【讨论】：

python 文档建议在处理大列表时使用 sorted() 而不是 heapq.nlargest()，尽管它们没有说明“大列表”有多大。 docs.python.org/library/heapq.html
@Matt 我在文档中找不到这样的建议-但很可能是这种情况-我建议 OP 运行一些 timeit 测试以找出最有效的方法供他使用。
@Lattyware 在 docs.python.org/library/heapq.html 上它说“后两个函数对于较小的 n 值表现最好。对于较大的值，使用 sorted( ) 函数。另外，当 n==1 时，使用内置的 min() 和 max() 函数效率更高。"

【解决方案4】：

heapq 的另一种使用方式

heapq.nlargest(n, range(len(a)), key=a.__getitem__)

正如在其他地方评论的那样，除非 a 非常大并且 n<<len(a)，否则它不会胜过排序，因为排序是 Python 中相对较快的操作。然而最终缓慢的 O(n) 总是会击败 O(n*log(n))

【讨论】：

是的，你是对的，对于大 n，慢 O(n) 比 O(n*log(n)) 慢，但是 heapq 模块已经智能实现，注意 n 值传递给 @987654328 @ 函数仅在 n 相对较小时才使用堆，并在 n 明显较大且趋于 sizeof 列表时切换到排序。
nlargest from 至少 v2.7.11 无论如何都会使用堆，但 v3.5.2 的行为与您描述的@sinister 一样
其实在 v3.5.2 中，sorted 只在 n 大于可迭代 @sinister 的大小时才使用