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

如果你想找到n最大元素的索引，bottleneck你可以使用 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

每个循环的平均时间：

• 瓶颈.partsort()：59 毫秒
• np.percentile()：54 毫秒

Answer 7

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

import heapq
heapq.nlargest(N, a)

获得N最大的成员。