我正在寻找一种快速计算滚动总和的方法,可能使用Numpy。这是我的第一个方法:
def func1(M, w):
Rtn = np.zeros((M.shape[0], M.shape[1]-w+1))
for i in range(M.shape[1]-w+1):
Rtn[:,i] = np.sum(M[:, i:w+i], axis=1)
return Rtn
M = np.array([[0., 0., 0., 0., 0., 1., 1., 0., 1., 1., 1., 0., 0.],
[0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 1., 1.],
[1., 1., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0.]])
window_size = 4
print func1(M, window_size)
[[ 0. 0. 1. 2. 2. 3. 3. 3. 3. 2.]
[ 1. 2. 2. 1. 1. 0. 0. 0. 1. 2.]
[ 3. 2. 1. 1. 1. 1. 1. 1. 0. 0.]]
我想阻止窗口(/ sum)在循环中重做,希望能让它快得多,所以我想出了以下函数,它将总和限制为滚动窗口的第一个和最后一个元素: / p>
def func2(M, w):
output = np.zeros((M.shape[0], M.shape[1]-w+1))
sum = np.sum(M[:, 0:w], axis=1)
output[:,0] = sum
for i in range(w, M.shape[1]):
sum = sum + M[:,i]- M[:,i-w]
output[:,i-w+1] = sum
return output
但令我惊讶的是,func2几乎比func1快得多:
In [251]:
M = np.random.randint(2, size=3000).reshape(3, 1000)
window_size = 100
%timeit func1(M, window_size)
10 loops, best of 3: 20.9 ms per loop
In [252]:
%timeit func2(M, w)
10 loops, best of 3: 15.5 ms per loop
我在这里遗漏了什么吗?你们知道更好吗,我的意思是更快地实现这个目标吗?
答案 0 :(得分:8)
改编自@ Jaime的答案:https://stackoverflow.com/a/14314054/553404
import numpy as np
def rolling_sum(a, n=4) :
ret = np.cumsum(a, axis=1, dtype=float)
ret[:, n:] = ret[:, n:] - ret[:, :-n]
return ret[:, n - 1:]
M = np.array([[0., 0., 0., 0., 0., 1., 1., 0., 1., 1., 1., 0., 0.],
[0., 0., 1., 0., 1., 0., 0., 0., 0., 0., 0., 1., 1.],
[1., 1., 0., 1., 0., 0., 0., 1., 0., 0., 0., 0., 0.]])
print(rolling_sum(M))
输出
[[ 0. 0. 1. 2. 2. 3. 3. 3. 3. 2.]
[ 1. 2. 2. 1. 1. 0. 0. 0. 1. 2.]
[ 3. 2. 1. 1. 1. 1. 1. 1. 0. 0.]]
计时
In [7]: %timeit rolling_sum(M, 4)
100000 loops, best of 3: 7.89 µs per loop
In [8]: %timeit func1(M, 4)
10000 loops, best of 3: 70.4 µs per loop
In [9]: %timeit func2(M, 4)
10000 loops, best of 3: 54.1 µs per loop