python - 将 2 个暗淡的 numpy 数组与向量相乘得到 3 个暗淡的数组

Question

假设我有一个二维数组和一个一维数组

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

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

我想要实现的是获得一个 3d 数组 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]])

有没有更简单的版本？

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更快。