Logo Questions Linux Laravel Mysql Ubuntu Git Menu
 

Numpy: Conditional sum rows using an mask array

Suppose I have array a and array mask,

Array a:
[[1,1,2]
 [2,2,3]
 [3,5,2]
 [2,3,4]]

Array mask:
[[0,1]
 [1,1]
 [1,0]
 [0,0]]

How can I generate a array c with shape(2,3) in a numpy way like below:

Array c:
[[5, 3],
 [7, 3],
 [5, 5]])

where the 1st column is the sum of rows of array a indicated by mask[:,0] and 2nd column is the sum of rows indicated by mask[:,1], like below:

c[:0] = a[1]+a[2]

c[:1] = a[0]+a[1]
like image 393
Allen Avatar asked Jul 02 '26 01:07

Allen


1 Answers

You can use the numpy dot product, which is essentially a matrix product, as documented here:

For 2-D arrays it is the matrix product

import numpy as np
np.dot(a.transpose(), mask)

# array([[5, 3],
#        [7, 3],
#        [5, 5]])
like image 125
Psidom Avatar answered Jul 03 '26 14:07

Psidom