广播 - 将(N×N)阵列乘以(M)阵列以具有(NxNxM)阵列

时间:2016-06-01 15:03:50

标签: python numpy

我有两个numpy ndarrays

a.shape # -> (3, 3)
b.shape # -> (5,)

我想计算一个由{/ p>定义的c数组

c.shape # -> (3, 3, 5)
c[...,i] = a*b[i]

c的确切形状并不重要,因为我总能调换它的轴。)

在我看来,到目前为止发布我的编码尝试会不必要地让人感到尴尬,就像a[:,None,:]*b[None,:]上的许多变化一样,这就足够了,不是吗?

我怀疑np.einsum()可能是答案,但其下标命令的语法超出了我的想法......

2 个答案:

答案 0 :(得分:5)

一个很好的解决方案是使用精彩的einsum函数:

>>> a = np.array([[1,2,3],[4,5,6],[7,8,9]])
>>> b = np.array([10, 100, 1000, 10000, 100000])
>>> c = np.einsum('ij,k->ijk', a, b)
>>> c[1, 2, 3]
60000
>>> a[1, 2]
6
>>> b[3]
10000

我喜欢它,因为它使变换ij,k -> ijk非常明确。也可以,更简洁的是使用简单的广播:

>>> c = a[:, :, None] * b[None, None, :]
>>> c.shape
(3L, 3L, 5L)

使用:使用维度None沿着广播进行广播的位置。

答案 1 :(得分:0)

另一种解决方案是使用np.outer,因为基本上我们在那里执行外部产品。现在,np.outer期望输入为1D数组,并且作为预处理步骤,它在执行元素乘法之前使输入变平。因此,np.outer(a,b)的输出为2D,我们需要将其重塑为所需的3D形状。因此,最终的实现看起来像这样 -

np.outer(a,b).reshape(a.shape+(-1,))

更明确的重塑方式就是这样 -

np.outer(a,b).reshape(a.shape+b.shape)