在 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

我遇到了这个问题，因为这个问题已经 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 的解决方案）并列。 但是，使用 np.percentile (numpy_2) 您已经对那些 topN 元素进行了排序，而其他解决方案则并非如此。另一方面，如果您也对这些元素的索引感兴趣，则百分位数没有用处。

我也添加了 pandas，它在下面使用了瓶颈，如果有的话 (http://pandas.pydata.org/pandas-docs/stable/install.html#recommended-dependencies)。如果您已经有一个 pandas Series 或 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 4

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

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

如果您想使用 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 6

你也可以使用 numpy 的百分位函数。就我而言，它比bottleneck.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(M-N,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 毫秒
• np.percentile(): 54 毫秒

Answer 7

如果将数组存储为数字列表没有问题，您可以使用

import heapq
heapq.nlargest(N, a)

获得N 最大的成员。