将函数应用于其索引的numpy矩阵的每一行

Question

我有一个形状为[n,m] 的numpy矩阵A和一个长度为n 的数组b.我需要的是采取的总和b [I]的第i行的至少元件甲.所以代码可能如下所示:

A = np.array([[1,2,3],
              [4,5,6],
              [7,8,9]])

b = np.array([2,3,1])

sums = magic_function() #sums = [3, 15, 7]

我考虑过np.apply_along_axis()函数,但在这种情况下,你的函数似乎只能依赖于行本身.

Answer 1

矢量化方法利用NumPy broadcasting每行创建有效的掩码,然后执行sum-reduction-

mask = b[:,None] > np.arange(A.shape[1])
out = (A*mask).sum(1)

或者,np.einsum得到reduction-

out = np.einsum('ij,ij->i',A,mask)

我们也可以用np.matmul/@ notation on Python 3.x-

out = (A[:,None] @ mask[...,None]).squeeze()