我想编写一个函数,它将一个flattened数组作为输入,并返回一个等长的数组,其中包含输入数组中前n个元素的总和,以及输出数组的初始n - 1
元素到NaN
。
例如,如果数组有十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
是否有更好/更有效/更多" Pythonic"这样做的方法?
提前感谢您的帮助。
答案 0 :(得分:9)
您可以使用np.cumsum
,并获取cumsum
ed数组的差异及其移位版本:
n = 3
arr = np.array([2, 4, 3, 7, 6, 1, 9, 4, 6, 5])
sum_arr = arr.cumsum()
shifted_sum_arr = np.concatenate([[np.NaN]*(n-1), [0], sum_arr[:-n]])
sum_arr
=> array([ 2, 6, 9, 16, 22, 23, 32, 36, 42, 47])
shifted_sum_arr
=> array([ nan, nan, 0., 2., 6., 9., 16., 22., 23., 32.])
sum_arr - shifted_sum_arr
=> array([ nan, nan, 9., 14., 16., 14., 16., 14., 19., 15.])
IMO,这是一种更加 numpyish 的方法,主要是因为它避免了循环。
<强>计时强>
def cumsum_app(flat_array, n):
sum_arr = flat_array.cumsum()
shifted_sum_arr = np.concatenate([[np.NaN]*(n-1), [0], sum_arr[:-n]])
return sum_arr - shifted_sum_arr
flat_array = np.random.randint(0,9,(100000))
%timeit cumsum_app(flat_array,10)
1000 loops, best of 3: 985 us per loop
%timeit cumsum_app(flat_array,100)
1000 loops, best of 3: 963 us per loop
答案 1 :(得分:7)
您基本上在那里执行1D convolution
,因此您可以使用np.convolve
,就像这样 -
# Get the valid sliding summations with 1D convolution
vals = np.convolve(flat_array,np.ones(n),mode='valid')
# Pad with NaNs at the start if needed
out = np.pad(vals,(n-1,0),'constant',constant_values=(np.nan))
示例运行 -
In [110]: flat_array
Out[110]: array([2, 4, 3, 7, 6, 1, 9, 4, 6, 5])
In [111]: n = 3
In [112]: vals = np.convolve(flat_array,np.ones(n),mode='valid')
...: out = np.pad(vals,(n-1,0),'constant',constant_values=(np.nan))
...:
In [113]: vals
Out[113]: array([ 9., 14., 16., 14., 16., 14., 19., 15.])
In [114]: out
Out[114]: array([ nan, nan, 9., 14., 16., 14., 16., 14., 19., 15.])
对于1D卷积,也可以使用Scipy's implementation
。对于大窗口大小,使用Scipy版本的运行时似乎更好,因为下面列出的运行时测试也会尝试调查。获取vals
的Scipy版本为 -
from scipy import signal
vals = signal.convolve(flat_array,np.ones(n),mode='valid')
NaNs
填充操作可以替换为np.hstack
:np.hstack(([np.nan]*(n-1),vals))
以获得更好的效果。
运行时测试 -
In [238]: def original_app(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
...:
...: def vectorized_app1(flat_array,n):
...: vals = np.convolve(flat_array,np.ones(n),mode='valid')
...: return np.hstack(([np.nan]*(n-1),vals))
...:
...: def vectorized_app2(flat_array,n):
...: vals = signal.convolve(flat_array,np.ones(3),mode='valid')
...: return np.hstack(([np.nan]*(n-1),vals))
...:
In [239]: flat_array = np.random.randint(0,9,(100000))
In [240]: %timeit original_app(flat_array,10)
1 loops, best of 3: 833 ms per loop
In [241]: %timeit vectorized_app1(flat_array,10)
1000 loops, best of 3: 1.96 ms per loop
In [242]: %timeit vectorized_app2(flat_array,10)
100 loops, best of 3: 13.1 ms per loop
In [243]: %timeit original_app(flat_array,100)
1 loops, best of 3: 836 ms per loop
In [244]: %timeit vectorized_app1(flat_array,100)
100 loops, best of 3: 16.5 ms per loop
In [245]: %timeit vectorized_app2(flat_array,100)
100 loops, best of 3: 13.1 ms per loop
答案 2 :(得分:4)
这里的其他答案可能更接近于您在速度和内存方面的要求,但为了完整性,您还可以使用列表理解来构建阵列:
a = np.array([2, 4, 3, 7, 6, 1, 9, 4, 6, 5])
N, n = a.shape[0], 3
np.array([np.NaN]*(n-1) + [np.sum(a[j:j+n]) for j in range(N-n+1)])
返回:
array([ nan, nan, 9., 14., 16., 14., 16., 14., 19., 15.])