时间戳

首页 > 解决方案 > 计算向量的总和

numpy - 计算向量的总和

问题描述

我有两个矩阵a，b并且想将它们之间的所有总和计算为张量。我怎样才能比执行以下代码更有效地做到这一点：

a = np.array([[1,2],[3,4],[5,6]])
b = np.array([[4,5],[6,7]])

n1 = a.shape[0]
n2 = b.shape[0]
f = a.shape[1]
c = np.zeros((n1,n2,f))

c = np.zeros((n1,n2,f))
for i in range(n1):
    for j in range(n2):
        c[i,j,:] = a[i,:] + b[j,:]

einstein-sum 之类的显然不起作用，外积也没有——有合适的方法吗？

标签： numpy

解决方案

您可以将循环表达式转换为广播表达式：

c[i,j,:] = a[i,:] + b[j,:]
c[i,j,:] = a[i,None,:] + b[None,j,:]  # fill in the missing dimensions
c = a[:,None,:] + b[None,:,:]


In [167]: a[:,None,:]+b[None,:,:]                                                                              
Out[167]: 
array([[[ 5,  7],
        [ 7,  9]],

       [[ 7,  9],
        [ 9, 11]],

       [[ 9, 11],
        [11, 13]]])
In [168]: _.shape                                                                                            
Out[168]: (3, 2, 2)

a[:,None]+b做同样的事情，因为前导None( np.newaxis) 是自动的，并且尾随:也是。

推荐阅读