如何计算一个二维数组的所有列之和(有效)

假设我有一个由四行三列组成的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])之和的1D 数组的有效方法是什么？在不需要遍历所有列的情况下，是否可以做到这一点？

224420 次浏览
小开
最佳答案

查看 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参数。

来自 初步麻木教程:

许多一元运算，例如计算所有元素的和 在数组中，作为 ndarray类的方法实现 默认情况下，这些操作应用于数组，就像它是一个列表一样 无论数字的形状如何，只要指定 axis 参数的指定轴应用操作 数组:

小开

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

sum = a.sum(axis=0)
小开

使用 axis参数:

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

然后 NumPy sum函数接受一个可选的轴参数，该参数指定希望执行的和沿着哪个轴:

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

或者，等价地说:

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

对列求和的其他替代方法包括

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

还有

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

如果行和列的数量在同一个数量级中，那么所有的可能性都大致相同:

enter image description here

但是，如果只有几列，那么 einsum和 dot解决方案的性能都明显优于 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.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)',
)
小开 
a.sum(0)

应该能解决问题。它是一个2dnp.array，你会得到所有列的和。axis=0是向下的维度，axis=1是向右的维度。


提交答案