Question

假设我有一个ndy数组，它是2d和1d数组

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] and C[:, :, 1] = A*B[1]的等价物，我们需要扩展A -

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

如果此处没有减少额外费用，基于broadcasting的会比einsum更快。

用向量乘以2dim numpy数组得到3dim数组

1 个答案: