python之有效地计算数组中 N 个最小数字的总和

我有一个代码，首先我需要对值进行排序，然后我需要对前 10 个元素求和。我很想使用 Numba 包来加快运行时间，但它不起作用，Numba 使代码比 Numpy 慢。

我的第一个测试，只是求和:

import numpy as np 
import numba 
np.random.seed(0) 
 
def SumNumpy(x): 
    return np.sum(x[:10]) 
 
@numba.jit() 
def SumNumpyNumba(x): 
    return np.sum(x[:10]) 

我的测试:

x = np.random.rand(1000000000) 
%timeit SumNumpy(x) 
%timeit SumNumpyNumba(x) 

结果:

100000 次循环，最好的 3 次:每次循环 6.8 µs

1000000 次循环，3 次最佳:每次循环 715 ns

没关系，Numba 做得很好。但是当我一起尝试 np.sort 和 np.sum 时:

def sumSortNumpy(x): 
    y = np.sort(x) 
    return np.sum(y[:10]) 
 
@numba.jit() 
def sumSortNumpyNumba(x): 
    y = np.sort(x) 
    return np.sum(y[:10]) 

和测试:

x = np.random.rand(100000) 
%timeit sumSortNumpy(x) 
%timeit sumSortNumpyNumba(x) 

结果:

100 次循环，最好的 3 次:每次循环 14.6 毫秒

10 个循环，3 个循环中的最佳:每个循环 20.6 毫秒

Numba/Numpy 比 Numpy 更慢。所以我的问题是我们是否可以改进函数“sumSortNumpyNumba”？

感谢帮助。

谢谢。

请您参考如下方法:

我们在排序后求和，因此前 N=10 元素的顺序无关紧要。因此，我们可以使用 np.argpartition 来避免排序步骤，并简单地为我们提供第一个 N 最小数字的组，这些数字可以在以后求和，就像这样 -

def sumSortNumPyArgpartition(x, N=10): 
    return x[np.argpartition(x, N)[:N]].sum() 

各种数据集的时间 -

In [39]: np.random.seed(0) 
    ...: x = np.random.rand(1000000) 
 
In [40]: %timeit sumSortNumpy(x) 
    ...: %timeit sumSortNumPyArgpartition(x) 
10 loops, best of 3: 78.6 ms per loop 
100 loops, best of 3: 12.3 ms per loop 
 
In [41]: np.random.seed(0) 
    ...: x = np.random.rand(10000000) 
 
In [42]: %timeit sumSortNumpy(x) 
    ...: %timeit sumSortNumPyArgpartition(x) 
1 loop, best of 3: 920 ms per loop 
10 loops, best of 3: 153 ms per loop 
 
In [43]: np.random.seed(0) 
    ...: x = np.random.rand(100000000) 
 
In [44]: %timeit sumSortNumpy(x) 
    ...: %timeit sumSortNumPyArgpartition(x) 
1 loop, best of 3: 10.6 s per loop 
1 loop, best of 3: 978 ms per loop