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

时间:2012-11-26 14:55:05

标签: python numpy

假设我有以下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])的有效方法是什么?这可以在不需要遍历所有列的情况下完成吗?

6 个答案:

答案 0 :(得分:87)

查看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.meannumpy.cumsumnumpy.std,也可以使用axis参数。

来自Tentative Numpy Tutorial

  

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

答案 1 :(得分:4)

使用numpy.sum。对于你的情况,它是

sum = a.sum(axis=0)

答案 2 :(得分:3)

使用axis参数:

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

答案 3 :(得分:3)

用于求和列的其他替代方法是

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

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

如果行数和列数处于同一数量级,则所有可能性大致相同:

enter image description here

但是,如果只有几列,则einsumdot解决方案明显优于numpy的sum(请注意对数刻度):

enter image description here

重现图表的代码:

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.show(
    # 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)',
    )

答案 4 :(得分:2)

然后NumPy sum函数接受一个可选的轴参数,该参数指定您希望执行求和的轴:

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

或等同地:

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

答案 5 :(得分:0)

a.sum(0)

应该解决问题。它是2d np.array,您将获得所有列的总和。 axis=0是指向下方的尺寸,而axis=1是指向右侧的尺寸。