Вычислить сумму попарных сумм двух столбцов массива

Question

Я ищу способ избежать вложенных циклов в следующем фрагменте, где A а также B двумерные массивы, каждый из формы (m, n) с m, n быть произвольными положительными целыми числами:

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, как @Divakar упоминается в #approach 2

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