一个快速的方法来find一个numpy数组中最大的N个元素

我知道我可以做到这一点，如下所示：

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

然而，这是一个完整的sorting，是非常缓慢的。

我不知道numpy是否提供了一些快速的方法。

Answer 1

bottleneck模块有一个快速的部分sorting方法，直接与Numpy数组一起工作： bottleneck.partition() 。

请注意， bottleneck.partition()返回sorting的实际值，如果你想sorting值的索引（what numpy.argsort()返回），你应该使用bottleneck.argpartition() 。

我做了基准testing：

• 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毫秒

Answer 2

numpy 1.8实现了与O（n）* log（n））完全sorting相对的O（n）时间执行部分sorting的partition和argpartition 。

 import numpy as np test = np.array([9,1,3,4,8,7,2,5,6,0]) temp = np.argpartition(-test, 4) result_args = temp[:4] temp = np.partition(-test, 4) result = -temp[:4] 

结果：

 >>> result_args array([0, 4, 8, 5]) # indices of highest vals >>> result array([9, 8, 6, 7]) # highest vals 

定时：

 In [16]: a = np.arange(10000) In [17]: np.random.shuffle(a) In [18]: %timeit np.argsort(a) 1000 loops, best of 3: 1.02 ms per loop In [19]: %timeit np.argpartition(a, 100) 10000 loops, best of 3: 139 us per loop In [20]: %timeit np.argpartition(a, 1000) 10000 loops, best of 3: 141 us per loop 

Answer 3

拟议的瓶颈解决scheme中的每个负号

 -bottleneck.partsort(-a, 10)[:10] 

制作数据的副本。 我们可以通过做删除副本

 bottleneck.partsort(a, a.size-10)[-10:] 

另外提出的numpy解决scheme

 a.argsort()[-10:] 

返回索引不是值。 解决方法是使用索引来查找值：

 a[a.argsort()[-10:]] 

两种瓶颈解决scheme的相对速度取决于初始arrays中元素的sorting，因为这两种方法在不同点处分割数据。

换句话说，任何一个特定的随机数组的定时可以使任何一种方法看起来更快。

对100个随机数组进行平均，每个数组有1000000个元素

 -bn.partsort(-a, 10)[:10]: 1.76 ms per loop bn.partsort(a, a.size-10)[-10:]: 0.92 ms per loop a[a.argsort()[-10:]]: 15.34 ms per loop 

时间码如下：

 import time import numpy as np import bottleneck as bn def bottleneck_1(a): return -bn.partsort(-a, 10)[:10] def bottleneck_2(a): return bn.partsort(a, a.size-10)[-10:] def numpy(a): return a[a.argsort()[-10:]] def do_nothing(a): return a def benchmark(func, size=1000000, ntimes=100): t1 = time.time() for n in range(ntimes): a = np.random.rand(size) func(a) t2 = time.time() ms_per_loop = 1000000 * (t2 - t1) / size return ms_per_loop t1 = benchmark(bottleneck_1) t2 = benchmark(bottleneck_2) t3 = benchmark(numpy) t4 = benchmark(do_nothing) print "-bn.partsort(-a, 10)[:10]: %0.2f ms per loop" % (t1 - t4) print "bn.partsort(a, a.size-10)[-10:]: %0.2f ms per loop" % (t2 - t4) print "a[a.argsort()[-10:]]: %0.2f ms per loop" % (t3 - t4) 

Answer 4

也许heapq.nlargest

 import numpy as np import heapq x = np.array([1,-5,4,6,-3,3]) z = heapq.nlargest(3,x) 

结果：

 >>> z [6, 4, 3] 

如果要查找使用bottleneck的n最大元素的索引，可以使用bottleneck.argpartsort

 >>> x = np.array([1,-5,4,6,-3,3]) >>> z = bottleneck.argpartsort(-x, 3)[:3] >>> z array([3, 2, 5] 

Answer 5

你也可以使用numpy的百分位数function。 在我的情况下，它稍微快瓶颈.partsort（）：

 import timeit import bottleneck as bn N,M,K = 10,1000000,100 start = timeit.default_timer() for k in range(K): a=np.random.uniform(size=M) tmp=-bn.partsort(-a, N)[:N] stop = timeit.default_timer() print (stop - start)/K start = timeit.default_timer() perc = (np.arange(MN,M)+1.0)/M*100 for k in range(K): a=np.random.uniform(size=M) tmp=np.percentile(a,perc) stop = timeit.default_timer() print (stop - start)/K 

每个循环的平均时间：

• bottleneck.partsort（）：59 ms
• np.percentile（）：54毫秒

Answer 6

我有这个问题，因为这个问题是5岁，我不得不重做所有的基准，并改变瓶颈的语法（没有partsort ，现在是partition ）。

我用与kwgoodman相同的论点，除了检索的元素数量，我增加到50（为了更好地适应我的具体情况）。

我得到了这些结果：

 bottleneck 1: 01.12 ms per loop bottleneck 2: 00.95 ms per loop pandas : 01.65 ms per loop heapq : 08.61 ms per loop numpy : 12.37 ms per loop numpy 2 : 00.95 ms per loop 

所以，bottleneck_2和numpy_2（adas的解决scheme）是捆绑在一起的。 但是，使用np.percentile （numpy_2），您已经sorting了这些topN元素，而其他解决scheme则不是这样。 另一方面，如果你也对这些元素的索引感兴趣，百分位数是没有用的。

我也添加了pandas，它使用下面的瓶颈，如果可用（ http://pandas.pydata.org/pandas-docs/stable/install.html#recommended-dependencies ）。 如果你已经有了一个pandas系列或数据框，那么你已经掌握了一切，只需要使用nlargest ，就完成了。

用于基准的代码如下（python 3，请）：

 import time import numpy as np import bottleneck as bn import pandas as pd import heapq def bottleneck_1(a, n): return -bn.partition(-a, n)[:n] def bottleneck_2(a, n): return bn.partition(a, a.size-n)[-n:] def numpy(a, n): return a[a.argsort()[-n:]] def numpy_2(a, n): M = a.shape[0] perc = (np.arange(Mn,M)+1.0)/M*100 return np.percentile(a,perc) def pandas(a, n): return pd.Series(a).nlargest(n) def hpq(a, n): return heapq.nlargest(n, a) def do_nothing(a, n): return a[:n] def benchmark(func, size=1000000, ntimes=100, topn=50): t1 = time.time() for n in range(ntimes): a = np.random.rand(size) func(a, topn) t2 = time.time() ms_per_loop = 1000000 * (t2 - t1) / size return ms_per_loop t1 = benchmark(bottleneck_1) t2 = benchmark(bottleneck_2) t3 = benchmark(pandas) t4 = benchmark(hpq) t5 = benchmark(numpy) t6 = benchmark(numpy_2) t0 = benchmark(do_nothing) print("bottleneck 1: {:05.2f} ms per loop".format(t1 - t0)) print("bottleneck 2: {:05.2f} ms per loop".format(t2 - t0)) print("pandas : {:05.2f} ms per loop".format(t3 - t0)) print("heapq : {:05.2f} ms per loop".format(t4 - t0)) print("numpy : {:05.2f} ms per loop".format(t5 - t0)) print("numpy 2 : {:05.2f} ms per loop".format(t6 - t0)) 

Answer 7

如果将数组存储为数字列表并不成问题，则可以使用

 import heapq heapq.nlargest(N, a) 

得到N最大的成员。