我想写一个函数,它以一个展平的数组作为输入,并返回一个等长的数组,其中包含输入数组中前n个元素的和,输出数组的初始n - 1
元素设置为NaN
。
例如,如果数组有10个elements = [2, 4, 3, 7, 6, 1, 9, 4, 6, 5]
和n = 3
,那么结果数组应该是[NaN, NaN, 9, 14, 16, 14, 16, 14, 19, 15]
。
我想到的一个方法是:
def sum_n_values(flat_array, n):
sums = np.full(flat_array.shape, np.NaN)
for i in range(n - 1, flat_array.shape[0]):
sums[i] = np.sum(flat_array[i - n + 1:i + 1])
return sums
有没有更好/更有效/更“Python”的方法来做到这一点?
3条答案
按热度按时间sf6xfgos1#
你可以使用
np.cumsum
,并取cumsum
ed数组和它的移位版本的差:IMO,这是一个更 numpyish 的方式来做到这一点,主要是因为它避免了循环。
时间
aelbi1ox2#
你基本上是在执行
1D convolution
,所以你可以使用np.convolve
,就像这样-样品运行-
对于1D卷积,也可以使用
Scipy's implementation
。Scipy版本的运行时似乎更适合大的窗口大小,因为接下来列出的运行时测试也将尝试调查。获取vals
的Scipy版本是-NaNs
填充操作可以被np.hstack
替换:np.hstack(([np.nan]*(n-1),vals))
以获得更好的性能。测试-
5w9g7ksd3#
这里的其他答案可能更接近你在速度和内存方面所寻求的,但为了完整起见,你也可以使用列表解析来构建你的数组:
返回: