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

Question

我有两个numpy人的ndarray

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

我想要计算一个c由...定义的数组

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

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

在我看来,到目前为止发布我的编码尝试会不必要地令人尴尬,它就像许多变种一样,a[:,None,:]*b[None,:]而且足够说,不是吗？

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

Answer 1

一个很好的解决方案是使用梦幻般的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沿着一个广播.