Numpy:使用掩码数组的条件求和行

All*_*len 3 python arrays numpy

假设我有数组a和数组掩码,

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

Array mask:
[[0,1]
 [1,1]
 [1,0]
 [0,0]]
Run Code Online (Sandbox Code Playgroud)

如何以如下所示的numpy方式生成具有形状(2,3)的数组c:

Array c:
[[5, 3],
 [7, 3],
 [5, 5]])
Run Code Online (Sandbox Code Playgroud)

其中第1列是由掩码[:,0]表示的数组a的总和,第2列是掩码[:,1]指示的行的总和,如下所示:

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

c[:1] = a[0]+a[1]
Run Code Online (Sandbox Code Playgroud)

Psi*_*dom 5

您可以使用numpy点积,这基本上是一个矩阵的产品,如记录在这里:

对于二维阵列,它是矩阵产品

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

# array([[5, 3],
#        [7, 3],
#        [5, 5]])
Run Code Online (Sandbox Code Playgroud)