Python/NumPy:实现运行总和(但不完全)

Question

给定两个相等长度的数组,一个保持数据,一个保持结果但最初设置为零,例如:

a = numpy.array([1, 0, 0, 1, 0, 1, 0, 0, 1, 1])
b = numpy.array([0, 0, 0, 0, 0, 0, 0, 0, 0, 0])

我想计算a中三个相邻元素的所有可能子集的总和.如果和为0或1,则b中的三个对应元素保持不变; 只有当总和超过1时,b中的三个对应元素才被设置为1,因此在计算b变为之后

array([0, 0, 0, 1, 1, 1, 0, 1, 1, 1])

一个简单的循环将完成此任务:

for x in range(len(a)-2):
    if a[x:x+3].sum() > 1:
        b[x:x+3] = 1

在此之后,b具有所需的形式.

我必须为大量数据执行此操作,因此速度是一个问题.NumPy有更快的方式来执行上述操作吗？

(我知道这类似于卷积,但不完全相同).

Answer 1

您可以从卷积开始,选择超过1的值,最后使用"扩张":

b = numpy.convolve(a, [1, 1, 1], mode="same") > 1
b = b | numpy.r_[0, b[:-1]] | numpy.r_[b[1:], 0]

由于这避免了Python循环,它应该比你的方法更快,但我没有做时间.

另一种方法是使用第二个卷积来扩张:

kernel = [1, 1, 1]
b = numpy.convolve(a, kernel, mode="same") > 1
b = numpy.convolve(b, kernel, mode="same") > 0

如果您有可用的SciPy,那么扩张的另一个选择是

b = numpy.convolve(a, [1, 1, 1], mode="same") > 1
b = scipy.ndimage.morphology.binary_dilation(b)

编辑:通过做一些时间,我发现这个解决方案对于大型数组似乎是最快的:

b = numpy.convolve(a, kernel) > 1
b[:-1] |= b[1:]  # Shift and "smearing" to the *left* (smearing with b[1:] |= b[:-1] does not work)
b[:-1] |= b[1:]  # … and again!
b = b[:-2]

对于一百万个条目的数组,它比我的机器上的原始方法快200多倍.正如EOL在评论中指出的那样,这个解决方案可能被认为有点脆弱,因为它取决于NumPy的实现细节.