我想做一个 vectorized sum 使用 numpy 的数组 masked 指数

例如,如果没有遮罩:

import numpy as np

# data to be used in a vectorized sum operation
data = np.array([[1,0,0,0,0,0],
                 [0,1,0,0,0,0],
                 [0,0,1,0,0,0]])

# data indices i wish to sum together
idx = np.array([[0,1,2],   # sum data rows 0,1 and 2
                [2,1,1]])  # sum data rows 2,1 and 1

# without a mask this is straighforward
print np.sum(data[idx],axis=1)
#[[1 1 1 0 0 0]
# [0 2 1 0 0 0]]

现在使用掩码时,如果不在掩码索引数组上循环,我想不出如何执行此操作:

# introduce a mask
mask = np.array([[True,  True, True],  # sum data rows 0,1 and 2
                 [False, True, True]]) # sum data rows 1 and 1 (masking out idx[1,0])

summed = np.zeros((idx.shape[0],data.shape[1]),dtype='int')
for i in xrange(idx.shape[0]):
    summed[i] =  np.sum(data[idx[i][mask[i]]],axis=0)
print summed
#[[1 1 1 0 0 0]
 #[0 2 0 0 0 0]]

问题

有没有合适的方法来实现这种无回路的操作?