将numpy数组拆分为不同的部分,然后逐行求和

时间:2018-08-02 15:09:08

标签: numpy

我有一个数组列表

k=[array([0, 2], dtype=int64),
array([1], dtype=int64),
array([3, 4], dtype=int64)]

还有一个数组

e=array([[0., 4., 0., 0.],
         [1., 2., 2., 0.],
         [0., 0., 2., 4.],
         [0., 0., 0., 0.],
         [0., 0., 0., 3.]])

我想得到一个如下的数组x

array([[0,4,2,4],
       [1,2,2,0],
       [0,0,0,3]])

这里

 x[0] = e[0]+e[2]
 x[1] = e[1]
 x[2] = e[3]+e[4]

我希望以矢量化方式(无循环)进行操作。这可能吗?

1 个答案:

答案 0 :(得分:0)

如果输入数组的大小都相同,则可以使用完整的numpy解决方案:

k = np.array([[0, 2], [1, 1], [3, 4]], dtype=int)

e = np.array([[0., 4., 0., 0.],
         [1., 2., 2., 0.],
         [0., 0., 2., 4.],
         [0., 0., 0., 0.],
         [0., 0., 0., 3.]])

e[k].sum(axis=1)

如果没有,我认为有必要使用循环:

import numpy as np

k = [np.array([0, 2], dtype=int),
     np.array([1], dtype=int),
     np.array([3, 4], dtype=int)]

e = np.array([[0., 4., 0., 0.],
         [1., 2., 2., 0.],
         [0., 0., 2., 4.],
         [0., 0., 0., 0.],
         [0., 0., 0., 3.]])

np.array([ e[ki, :].sum(axis=0) for ki in k ])