在numpy中只对前两个值求和



如果我有一个像[1,2,3,4,1,2,3,4]这样的列表,我想对这些值求和:

[6, 9, 8, 6, 7, 4]

其中公式为value + value_one_ahead + value_two_ahead

该操作可以解释为输入与[1,1,1]的(离散)卷积,因此您可以使用numpy.convolve:

In [178]: x = np.array([1, 2, 3, 4, 1, 2, 3, 4])
In [179]: np.convolve(x, np.ones(3, dtype=x.dtype), mode='valid')
Out[179]: array([6, 9, 8, 7, 6, 9])

参见文档字符串mode参数的描述。如果可用的mode选项不能以您想要的方式处理输入的结束,您可能必须调整返回值或添加一两行额外的代码。例如,为了获得更接近您在问题中所写的内容,您可以使用mode='full'(这是默认值),然后丢弃前两个元素(通过获取切片[2:]):

In [181]: np.convolve(x, np.ones(3, dtype=x.dtype), mode='full')[2:]
Out[181]: array([6, 9, 8, 7, 6, 9, 7, 4])

注:np是用import numpy as np定义的。np.ones(3, dtype=x.dtype)创建一个长度和数据类型为1的数组:

In [180]: np.ones(3, dtype=x.dtype)
Out[180]: array([1, 1, 1])
l =  [1, 2, 3, 4, 1, 2, 3, 4]
result = []
for i in range(len(l)-2):
    result.append(
        l[i] + l[i+1] + l[i+2]
    )
print result

这是一个只有numpy的解决方案:

>>> arr[:-2] + arr[1:-1] + arr[2:]
array([6, 9, 8, 7, 6, 9])

最新更新