Search code examples
pythonarraysnumpyloopsrows

how to loop Numpy Array per row


I have two numpy Arrays

X = np.array([[0,1,0,1,1],
         [0,1,1,1,0]])

X2 = np.array([[0.2,0.5,0.1,0.5,0.5],
          [0.3,0.6,0.6,0.6,0.4]])

What I want is to get a new array with values of X2 where X is 0 and sum this per row so my output should be

[[0.3][0.7]]

I used

X2[X==0]

This gives me all the 0s from the whole array not per row so I get he sum of the whole array of 0s rather than the sum per row. Is there a slicing function or something I can use to just get the sums of the row?


Solution

  • You can use X as a mask, then multiply that mask with X2 and sum across axis 1, with keepdims=True:

    >>> np.sum((X==0) * X2, axis=1, keepdims=True)
    
    array([[0.3],
           [0.7]])