Summing up non-zero elements in numpy array

Question

I would like to sum up non-zero elements of an array (1-d) but do it separately for positive and negative integers (they can only be ones and two) and also display zeros where they are.

An example of an array:

array = np.array([0, 0, 0, -1, -1, 0, 1, 2, 1, 1, 0, -1, 0, 1, 1, -1, -2])

Output:

array([0, 0, 0, -2, 0, 5, 0, -1, 0, 2, -3])

I think my problem is that I have no idea how to separate the sequences of positive and negative values in array.

Answer 1

Here's one way -

def sum_by_signs(a):
    m1 = a<0
    m2 = a>0
    m0 = a==0 # or ~m1 & ~m2
    p = np.flatnonzero(m0[:-1] | np.diff(m1) | np.diff(m2))+1
    return np.add.reduceat(a, np.r_[0,p])

Or bring that np.r_[0 part into boolean construction part -

def sum_by_signs_v2(a):
    m1 = a<0
    m2 = a>0
    m0 = a==0 # or ~m1 & ~m2
    p = np.flatnonzero(np.r_[True, m0[:-1] | np.diff(m1) | np.diff(m2)])
    return np.add.reduceat(a, p)

Explanation

We are starting off the idea to split the array into "islands" based on sign changes or when we encounter a sequence of 0s, in which case we are looking to split each element as a separate one. By splitting, think of it as list of lists, if that makes it easier to understand. Now, the game is how do we get those splits. We need indices that signifies the start, stop indices of those islands. As said earlier, there are three cases, sign changes from + to - or vice versa or sequence of 0s.

Hence, that structure of boolean masks are used to give those indices with one-off slices to detect sign changes from + to - and vice versa with a combination of np.diff(m1) | np.diff(m2). Final one of m0[:-1] is for sequence of 0s. These indices are then fed to np.add.reduceat to get intervaled summations.

Sample runs -

In [208]: a
Out[208]: array([ 0,  0,  0, -1, -1,  0,  1,  2,  1,  1,  0, -1,  0,  1,  1, -1, -2])

In [209]: sum_by_signs(a)
Out[209]: array([ 0,  0,  0, -2,  0,  5,  0, -1,  0,  2, -3])

In [211]: a
Out[211]: array([ 1,  2,  0, -1, -1,  0,  1,  2,  1,  1,  0, -1,  0,  1,  1, -1, -2])

In [212]: sum_by_signs(a)
Out[212]: array([ 3,  0, -2,  0,  5,  0, -1,  0,  2, -3])

In [214]: a
Out[214]: 
array([ 1,  2,  0, -1, -1,  0,  1,  2,  1,  1,  0, -1,  0,  1,  1, -1, -2,
        0])

In [215]: sum_by_signs(a)
Out[215]: array([ 3,  0, -2,  0,  5,  0, -1,  0,  2, -3,  0])