NumPy:对每列n个矩阵求和

时间:2014-11-18 10:10:23

标签: python numpy matrix sum vectorization

我希望对矩阵的每n列求和。如何在不使用for循环的情况下以简单的方式完成此操作?这就是我现在所拥有的:

n = 3  #size of a block we need to sum over
total = 4  #total required sums
ncols = n*total
nrows = 10
x = np.array([np.arange(ncols)]*nrows)

result = np.empty((total,nrows))
for i in range(total):
    result[:,i] =  np.sum(x[:,n*i:n*(i+1)],axis=1)

结果将是

array([[  3.,  12.,  21.,  30.],
       [  3.,  12.,  21.,  30.],
        ...
       [  3.,  12.,  21.,  30.]])

如何对此操作进行矢量化?

1 个答案:

答案 0 :(得分:4)

这是一种方式;首先将x重新整形为3D数组,然后对最后一个轴求和:

>>> x.reshape(-1, 4, 3).sum(axis=2)
array([[ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30],
       [ 3, 12, 21, 30]])