假设我有一个距离数组x=[1,2,1,3,3,2,1,5,1,1]
。
我想要得到满足累积总和为10时的x数组下标,即idx=[4,9]。
因此,在符合条件后,累积总和会重新开始计算。
我可以使用循环来实现,但对于大型数组,循环速度较慢,我想知道是否可以以向量化的方式进行。
sumlm = np.frompyfunc(lambda a,b:a+b if a < 10 else b,2,1)
newx=sumlm.accumulate(x, dtype=np.object)
newx
array([1, 3, 4, 7, 10, 2, 3, 8, 9, 10], dtype=object)
np.nonzero(newx==10)
(array([4, 9]),)
np.random.seed([3, 1415]); x = np.random.randint(100, size=1_000_000).tolist()
进行检查。 - piRSquaredfrompyfunc
往往比更明确的 Python 循环快 2 倍。 - hpauljfrompyfunc
会生成一个带有像accumulate
这样的方法的ufunc
。 - hpaulj循环并不总是坏的(尤其是当你需要循环时)。另外,没有工具或算法可以比O(n)更快地完成此操作。因此,让我们编写一个好的循环。
def cumsum_breach(x, target):
total = 0
for i, y in enumerate(x):
total += y
if total >= target:
yield i
total = 0
list(cumsum_breach(x, 10))
[4, 9]
Numba是一个需要安装的第三方库。
Numba对支持的特性非常挑剔,但它可以正常工作。
此外,正如Divakar所指出的那样,Numba在处理数组方面具有更好的性能。
from numba import njit
@njit
def cumsum_breach_numba(x, target):
total = 0
result = []
for i, y in enumerate(x):
total += y
if total >= target:
result.append(i)
total = 0
return result
cumsum_breach_numba(x, 10)
因为我想这么做 ¯\_(ツ)_/¯
设置
np.random.seed([3, 1415])
x0 = np.random.randint(100, size=1_000_000)
x1 = x0.tolist()
准确度
i0 = cumsum_breach_numba(x0, 200_000)
i1 = list(cumsum_breach(x1, 200_000))
assert i0 == i1
时间
%timeit cumsum_breach_numba(x0, 200_000)
%timeit list(cumsum_breach(x1, 200_000))
582 µs ± 40.2 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
64.3 ms ± 5.66 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
Numba的速度大约是快100倍。
为了进行更加公正的比较,我将一个列表转换成了Numpy数组。
%timeit cumsum_breach_numba(np.array(x1), 200_000)
%timeit list(cumsum_breach(x1, 200_000))
43.1 ms ± 202 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
62.8 ms ± 327 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
total > 10
时我们可以使用 break 来使它更快。 - Quang Hoang这里有一个与numba和数组初始化有关的例子 -
from numba import njit
@njit
def cumsum_breach_numba2(x, target, result):
total = 0
iterID = 0
for i,x_i in enumerate(x):
total += x_i
if total >= target:
result[iterID] = i
iterID += 1
total = 0
return iterID
def cumsum_breach_array_init(x, target):
x = np.asarray(x)
result = np.empty(len(x),dtype=np.uint64)
idx = cumsum_breach_numba2(x, target, result)
return result[:idx]
时序
包括@piRSquared的解决方案
,并使用同一篇帖子中的基准测试设置-
In [58]: np.random.seed([3, 1415])
...: x = np.random.randint(100, size=1000000).tolist()
# @piRSquared soln1
In [59]: %timeit list(cumsum_breach(x, 10))
10 loops, best of 3: 73.2 ms per loop
# @piRSquared soln2
In [60]: %timeit cumsum_breach_numba(np.asarray(x), 10)
10 loops, best of 3: 69.2 ms per loop
# From this post
In [61]: %timeit cumsum_breach_array_init(x, 10)
10 loops, best of 3: 39.1 ms per loop
Numba: 追加 vs 数组初始化
为了更加深入地研究数组初始化的帮助作用,这似乎是两个 Numba 实现之间的重大区别,让我们对数组数据进行计时,因为数组数据的创建本身对运行时间影响很大,它们都依赖于它 -
In [62]: x = np.array(x)
In [63]: %timeit cumsum_breach_numba(x, 10)# with appending
10 loops, best of 3: 31.5 ms per loop
In [64]: %timeit cumsum_breach_array_init(x, 10)
1000 loops, best of 3: 1.8 ms per loop
In [65]: %timeit cumsum_breach_array_init(x, 10).copy()
100 loops, best of 3: 2.67 ms per loop
result [:idx]
的副本,以避免泄漏 result [idx:]
的内存。 - Paul Panzer
cumsum
,可以尝试使用以下代码:x.cumsum()%10
。如果x
是一个数组,这个方法非常快速。但是这种方法可能不适合你,因为如何处理边缘情况还不清楚。例如,如果cumsum
等于11,它应该重置为0还是1呢? - Brenlla