计算两个数组列的成对和的总和

Question

我正在寻找一种方法来避免以下片段中的嵌套循环，其中A和B是二维数组，每个形状(m, n)与m，n beeing正数整数：

import numpy as np

m, n = 5, 2
a = randint(0, 10, (m, n))
b = randint(0, 10, (m, n))

out = np.empty((n, n))
for i in range(n):
    for j in range(n):
        out[i, j] = np.sum(A[:, i] + B[:, j])

上述逻辑大致相当于

np.einsum('ij,ik', A, B)

除了einsum计算产品的总和。

有没有一种方法，相当于einsum，计算一笔总和？或者我是否必须为此操作编写扩展名？

Answer 1

einsum需要执行元素乘法，然后它进行求和（可选）。 因此，它可能不适用于解决我们的案例。 继续阅读！

方法＃1

我们可以利用broadcasting使第一个轴对齐，第二个轴在将尺寸扩展到3D之后进行元素求和。最后，我们需要沿第一轴求和 -

(A[:,:,None] + B[:,None,:]).sum(0)

方法＃2

我们可以简单地做外部添加每个的列式总结 -

A.sum(0)[:,None] + B.sum(0)

方法＃3

因此，引入einsum -

np.einsum('ij->j',A)[:,None] + np.einsum('ij->j',B)

Answer 2

你也可以使用numpy.ufunc.outer，特别是在这里numpy.add.outer沿着0轴求和之后在#approach 2中提到的@Divakar

In [126]: numpy.add.outer(a.sum(0),  b.sum(0))
Out[126]: 
array([[54, 67],
       [43, 56]])