将2dim numpy数组与向量相乘得到3dim数组

2024-03-29 12:41:06 发布

您现在位置:Python中文网/ 问答频道 /正文

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

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]])

有更简单的版本吗?你知道吗


Tags: in版本numpynp数组outarrayij
1条回答
网友
1楼 · 发布于 2024-03-29 12:41:06

要使用B沿第一个轴缩放,我们可以简单地使用^{},就像这样-

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的方法快。你知道吗

相关问题 更多 >