一种快速查找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

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

Answer 3

建议的瓶颈解决方案中的每个负号

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

制作数据的副本.我们可以通过这样做来删除副本

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

也提出了numpy解决方案

a.argsort()[-10:]

返回索引而不是值.修复是使用索引来查找值:

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

两个瓶颈解决方案的相对速度取决于初始阵列中元素的排序,因为这两种方法在不同点分割数据.

换句话说,使用任何一个特定随机数组的时序可以使任一方法看起来更快.

平均100个随机数组的时序,每个数组包含1,000,000个元素

-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

我有这个问题,因为这个问题是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

因此,瓶颈_2和numpy_2(adas的解决方案)是并列的.但是,使用np.percentile(numpy_2)你已经排序了那些topN元素,而其他解决方案则不然.另一方面,如果您对这些元素的索引感兴趣,百分位数也没用.

我也添加了pandas,它使用了下面的瓶颈(如果有的话)(http://pandas.pydata.org/pandas-docs/stable/install.html#recommended-dependencies).如果你已经开始使用pandas系列或DataFrame,那么你就掌握得很好,只需使用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(M-n,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 5

也许 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]

如果你想找到你可以使用 的n最大元素的索引bottleneckbottleneck.argpartsort

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