python - 有效地计算数组中 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 毫秒

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