Fast rolling-sum for list of data vectors (2d matrix)

孤者浪人 提交于 2019-12-03 14:42:32
YXD

Adapted from @Jaime's answer here: 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)) 

Output

[[ 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.]]

Timings

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
易学教程内所有资源均来自网络或用户发布的内容,如有违反法律规定的内容欢迎反馈
该文章没有解决你所遇到的问题?点击提问,说说你的问题,让更多的人一起探讨吧!