有效地计算阵列中N个最小数字的总和

Question

我有一个代码，首先我需要对值进行排序，然后我需要总结前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 ms每个循环

Numba / Numpy变得比Numpy慢。所以我的问题是否有什么可以改善功能“sumSortNumpyNumba”？

我很感激帮助。

谢谢。

Answer 1

我们在排序后进行求和，因此顺序在第一个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