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

问题描述

我知道我可以这样做:

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.

推荐答案

瓶颈 模块有一个直接与 Numpy 数组一起工作的快速部分排序方法:bottleneck.partition().

注意 bottleneck.partition() 返回实际排序的值，如果你想要排序值的索引(什么 numpy.argsort() 返回)你应该使用 bottleneck.argpartition().

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().

我已经进行了基准测试:

I've benchmarked:

• 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屋！