How to sum values in a numpy 3D array with conditional?

Question

I have a Numpy 3d array with 0 and 1 values, something like this:

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

       [[ 1, 1, 1, 1],
        [ 0, 1, 0, 1]]])

I would like to "add" (+ operation) each value in a "line" of the array following a specific condition: If I have consecutive "1" values, I put the sum. If I have 0, I keep it as it is. After a "0" values I restart counting.

The output that I would like to get is:

array([[[ 2, 0, 1],
        [ 0, 0, 2]],

       [[ 4],
        [ 0, 1, 0, 1]]])

The output may be "lines" with different sizes. Can I still do it with numpy? I searched on the forums, numpy tools, but I didn't find anything regarding my specific problem. If someone could point me towards the proper documentation/tool I would appreciate it. Thank you.

Divakar · Accepted Answer · 2019-10-29 14:55:51Z

2

Here's one way -

def sum_groups(a):
    z = np.zeros(a.shape[:-1] + (1,), dtype=a.dtype)
    b = np.concatenate((z,a,z),axis=-1)

    c = b.ravel()
    count = np.diff(np.flatnonzero(c[:-1]!=c[1:]))

    m2 = c[1:]>c[:-1]
    c[1:][m2] = count[::2]

    m3 = c==0
    m3[1:][m2] = 1
    m4 = m3.reshape(b.shape)
    m4[...,0] = 0
    m4[...,-1] = 0
    v = c[m3]

    rc = m4.sum(2).ravel()
    out = np.split(v,rc[:-1].cumsum())
    return out

Sample run -

In [7]: a
Out[7]: 
array([[[1, 1, 0, 1],
        [0, 0, 1, 1]],

       [[1, 1, 1, 1],
        [0, 1, 0, 1]]])

In [8]: sum_groups(a)
Out[8]: [array([2, 0, 1]), array([0, 0, 2]), array([4]), array([0, 1, 0, 1])]

Another one with more usage of boolean arrays for efficiency purposes -

def sum_groups_v2(a):
    p1 = a==1
    z1 = np.zeros(a.shape[:-1] + (1,), dtype=bool)
    b1 = np.concatenate((z1,p1,z1),axis=-1)

    c1 = b1.ravel()
    count1 = np.diff(np.flatnonzero(c1[:-1]!=c1[1:]))
    m33 = np.r_[False,c1[:-1]<c1[1:]].reshape(b1.shape)
    pp = np.zeros(b1.shape, dtype=int)
    pp[m33] = count1[::2]
    m33[...,1:-1][~p1] = 1
    v2 = pp[m33]
    rc2 = m33.sum(2).ravel()
    out2 = np.split(v2,rc2[:-1].cumsum())
    return out2

And one with region-based labelling -

from scipy.ndimage import label

def sum_groups_v3(a):
    z = np.zeros(a.shape[:-1] + (1,), dtype=a.dtype)
    b = np.concatenate((z,a),axis=-1)

    c = b.ravel()

    l = label(c)[0]
    unq,idxs,count = np.unique(l,return_index=True, return_counts=True)

    l[c==1] = -1
    l[idxs] = count
    p = l.reshape(b.shape)[...,1:]
    out = [j[j>=0] for i in p for j in i]
    return out

edited Oct 29, 2019 at 14:55

answered Oct 29, 2019 at 14:01

Divakar

222k19 gold badges273 silver badges374 bronze badges

Sign up to request clarification or add additional context in comments.

1 Comment

Zeineb Over a year ago

Thank you so so so much! It uses advanced functions (comparing to my basic programming level) that I wouldn't have found by my self, in a short amount of time! :)

Collectives™ on Stack Overflow

How to sum values in a numpy 3D array with conditional?

1 Answer 1

1 Comment

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

1 Comment

Your Answer

Sign up or log in

Post as a guest

Related