小编典典

在numpy数组中查找最大N个元素的快速方法

python

我知道我可以像下面这样:

import numpy as np
N=10
a=np.arange(1,100,1)
np.argsort()[-N:]

但是,由于它做了完整的排序,所以它非常慢。

我想知道numpy是否提供一些可以快速完成的方法。


阅读 339

收藏
2020-12-20

共1个答案

小编典典

bottleneck模块具有一种快速的局部排序方法,可直接与Numpy数组配合使用:bottleneck.partition()

请注意,bottleneck.partition()返回的是已排序的实际值,如果要使用已排序的值的索引(numpy.argsort()返回值),则应使用bottleneck.argpartition()

我已经进行了基准测试:

  • z = -bottleneck.partition(-a, 10)[:10]
  • z = a.argsort()[-10:]
  • z = heapq.nlargest(10, a)

其中a是一个随机的1,000,000个元素的数组。

时间安排如下:

  • bottleneck.partition():每个循环25.6毫秒
  • np.argsort():每个循环198毫秒
  • heapq.nlargest():每个循环358毫秒
2020-12-20