在numpy数组中查找最大N个元素的快速方法 [英] A fast way to find the largest N elements in an numpy array

问题描述

我知道我可以做到以下几点:

I know I can do it like the following:

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

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

However, it is very slow since it did a full sort.

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

I wonder whether numpy provide some methods the do it fast.

推荐答案

bottleneck 该模块具有直接用于Numpy数组的快速局部排序方法: bottleneck.partition()

请注意，bottleneck.partition()返回已排序的实际值，如果要获得已排序值的索引(numpy.argsort()返回的内容)，则应使用

Note that bottleneck.partition() returns the actual values sorted, if you want the indexes of the sorted values (what numpy.argsort() returns) you should use bottleneck.argpartition().

我已经进行了基准测试

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

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

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

where a is a random 1,000,000-element array.

时间如下:

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

• bottleneck.partition(): 25.6 ms per loop
• np.argsort(): 198 ms per loop
• heapq.nlargest(): 358 ms per loop

这篇关于在numpy数组中查找最大N个元素的快速方法的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！