Python/NumPy: implementing a running sum (but not quite)

Question

Given are two arrays of equal length, one holding data, one holding the results but initially set to zero, e.g.:

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])

I'd like to compute the sum of all possible subsets of three adjacent elements in a. If the sum is 0 or 1, the three corresponding elements in b are left unchanged; only if the sum exceeds 1 are the three corresponding elements in b set to 1, so that after the computation b becomes

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

A simple loop will accomplish this:

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

After this, b has the desired form.

I have to do this for a large amount of data, so speed is an issue. Is there a faster way in NumPy to carry out the operation above?

(I understand this is similar to a convolution, but not quite the same).

Answer 1

You can start with a convolution, choose the values that exceed 1, and finally use a "dilation":

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

Since this avoids the Python loop, it should be faster than your approach, but I didn't do timings.

An alternative is to use a second convolution to dilate:

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

If you have SciPy available, yet another option for the dilation is

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

Edit: By doing some timings, I found that this solution seems to be fastest for large arrays:

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]

For an array of one million entries, it was more than 200 times faster than your original approach on my machine. As pointed out by EOL in the comments, this solution might be considered a bit fragile, though, since it depends on implementation details of NumPy.

Answer 2

You can calculate the "convolution" sums in an efficient way with:

>>> a0 = a[:-2]
>>> a1 = a[1:-1]
>>> a2 = a[2:]
>>> a_large_sum = a0 + a1 + a2 > 1

Updating b can then be done efficiently by writing something that means "at least one of the three neighboring a_large_sum values is True": you first extend you a_large_sum array back to the same number of elements as a (to the right, to the left and to the right, and then to the left):

>>> a_large_sum_0 = np.hstack([a_large_sum, [False, False]])
>>> a_large_sum_1 = np.hstack([[False], a_large_sum, [False]])
>>> a_large_sum_2 = np.hstack([[False, False], a_large_sum])

You then obtain b in an efficient way:

>>> b = a_large_sum_0 | a_large_sum_1 | a_large_sum_2

This gives the result that you obtain, but in a very efficient way, through a leveraging of NumPy internal fast loops.

PS: This approach is essentially the same as Sven's first solution, but is way more pedestrian than Sven's elegant code; it is as fast, however. Sven's second solution (double convolve()) is even more elegant, and it is twice as fast.