如何计算2D numpy数组的所有列的总和(有效)

Question

假设我有以下2D numpy数组,包括四行和三列:

>>> a = numpy.arange(12).reshape(4,3)
>>> print(a)
[[ 0  1  2]
 [ 3  4  5]
 [ 6  7  8]
 [ 9 10 11]]

什么是生成包含所有列的总和的一维数组的有效方法(如[18, 22, 26])？这可以在不需要遍历所有列的情况下完成吗？

Answer 1

查看文档numpy.sum,特别注意axis参数.总结列:

>>> import numpy as np
>>> a = np.arange(12).reshape(4,3)
>>> a.sum(axis=0)
array([18, 22, 26])

或者,总结行:

>>> a.sum(axis=1)
array([ 3, 12, 21, 30])

其他聚合函数一样numpy.mean,numpy.cumsum并且numpy.std,例如,也采取了axis参数.

来自暂定的Numpy教程:

许多一元操作,例如计算数组中所有元素的总和,都是作为ndarray类的方法实现的.默认情况下,这些操作适用于数组,就像它是一个数字列表一样,无论其形状如何.但是,通过指定axis 参数,您可以沿数组的指定轴应用操作:

Answer 2

对列求和的其他替代方法是

numpy.einsum('ij->j', a)

和

numpy.dot(a.T, numpy.ones(a.shape[0]))

如果行数和列数处于相同数量级，则所有可能性的速度大致相同：

但是，如果只有几列，则einsum和dot解决方案都显着优于 numpy sum（注意对数刻度）：

重现图的代码：

import numpy
import perfplot


def numpy_sum(a):
    return numpy.sum(a, axis=1)


def einsum(a):
    return numpy.einsum('ij->i', a)


def dot_ones(a):
    return numpy.dot(a, numpy.ones(a.shape[1]))


perfplot.save(
    "out1.png",
    # setup=lambda n: numpy.random.rand(n, n),
    setup=lambda n: numpy.random.rand(n, 3),
    n_range=[2**k for k in range(15)],
    kernels=[numpy_sum, einsum, dot_ones],
    logx=True,
    logy=True,
    xlabel='len(a)',
    )

Answer 3

使用axis参数：

>> numpy.sum(a, axis=0)
  array([18, 22, 26])