CODEWITHSUNDEEP
pythonnumpy
Noah Toyonaga
Noah Toyonaga

Reputation: 23

Sum over variable size windows defined by custom boolean array

I can solve this problem using a loop structure, so I am specifically wondering whether there is any numpy/vectorized method of accomplishing this.

I have two arrays of values of the same length, e.g.

a = [1,2,3,4,5,6]
b = [False, True, True, False, False, True]

I would like to sum all the elements in a which correspond to each window of "true" positions in b and add them to the immediately preceding "false" values (again in a). So to complete the example I gave I would like to get the output

c = [6,4,11]

In this case we get: 6 (sum of 1,2,3 from the sequence [False, True, True, in b), 4 (no "True" positions immediately following in b) , 11 (from False, True in b).

I realize that may be hard to follow, so please let me know if another example/more explanation would be helpful.

Upvotes: 2

Views: 62

Answers (1)

AGN Gazer
AGN Gazer

Reputation: 8378

I will piggyback on @Divakar's comment. It is a great idea to use np.add.reduceat(). However, it seems to me that @Divakar's treatment of indices is too simplistic for this problem and a more complicated analysis is required. I think the following may produce what you are looking for:

idx = np.insert(
    np.flatnonzero(
        np.ediff1d(np.pad(b.astype(np.int), 1, 'constant', constant_values=0)) == -1
    ), 0, 0
)
wsum = np.add.reduceat(np.append(a, 0), idx)[:-1]

Test

With

>>> a = np.arange(1, 11)
>>> b = np.array([True, True, False, True, False, True, True, False, False, True])

I get:

>>> print(wsum)
[ 3  7 18 27]

Upvotes: 1

Related Questions