将2维numpy数组与向量相乘以获得3维数组

Question

假设我有一个numpy数组，它是2D的，还有一个1D的数组。

In [127]: A = np.array([[1, 2],[3, 4]])

In [128]: B = np.array([10, 100])

我希望实现的目标是获取一个三维数组C，其中C[:, :, 0] = A*B[0]和C[:, :, 1] = A*B[1]。我可以通过np.einsum来实现，但似乎有点过头了。

In [129]: np.einsum('ij, k -> ijk', A, B)[:, :, 0]
Out[129]: 
array([[10, 20],
       [30, 40]])

In [130]: np.einsum('ij, k -> ijk', A, B)[:, :, 1]
Out[130]: 
array([[100, 200],
       [300, 400]])

有更简单的版本吗？

- math

---

B[:,None,None]*A 的翻译是：B[:,None,None]*A。 - Divakar

使用新的（第三个）维度进行广播。就像@Divakar所说的那样。 - sshashank124

C = np.einsum('ij, k -> ijk', A, B) 已经在您的代码中使用了，您觉得如何？我怀疑 A[...,None]*B 可以更快地得到与 C[:, :, 0] = A*B[0] and C[:, :, 1] = A*B[1] 等效的结果。 - Divakar

@Divakar 很酷，非常感谢。如果您发布它，我很乐意将其作为答案接受。我认为应该是 B[None, None, :] *A 才能等同于我的解决方案。 - math

Answer 1

要沿第一轴使用B进行缩放，我们可以简单地使用broadcasting，如下所示 -

B[:,None,None]*A # with einsum : np.einsum('ij, k -> kij', A, B)

为了得到与 C[:, :, 0] = A*B[0] 和 C[:, :, 1] = A*B[1] 等价的结果，我们需要扩展 A。

A[...,None]*B # with einsum : np.einsum('ij, k -> ijk', A, B)

在这里没有进行任何求和操作，基于的操作比基于einsum的操作更快。

- Divakar