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

Question

我知道我可以这样做：

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

然而，它是非常缓慢的，因为它做了一个完整的排序。

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

Answer 1

numpy 1.8实现partition和argpartition，它们执行部分排序（在O（n）时间内，而不是完全排序（O（n）*log（n））。

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 2

^{}模块有一个快速的部分排序方法，直接与Numpy数组一起工作：^{}。

请注意，bottleneck.partition()返回已排序的实际值，如果需要排序值的索引（返回的是什么），则应使用^{}。

我的基准是：

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

其中a是一个随机的1000000元素数组。

时间安排如下：

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

Answer 3

提出的瓶颈解决方案中的每个负号

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

复制数据。我们可以通过

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

也是建议的numpy解决方案

a.argsort()[-10:]

返回索引而不是值。修复方法是使用索引查找值：

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

两个瓶颈解决方案的相对速度取决于初始数组中元素的顺序，因为这两种方法在不同的点上分割数据。

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

平均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)