通过使用 np.einsum 将 2D 数组的行广播到 3D 数组的每个切片进行乘法

Question

我有数组 A 和 B:

>>> import numpy as np

>>> A = np.ones((3,3,2))

>>> B = np.array([
    [0,0],
    [1,1],
    [2,2],
])

我想将 B 的每一行乘以 A 的每一片，这样 B 的每一行都会在 A 的每一片上广播，即是：

>>> np.array([A_slice*B_row for A_slice, B_row in zip(A, B)])
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

我想要执行此操作的最有效方法，我认为可能是使用 np.einsum（但是，如果您认为使用其他方法更快，例如我在下面提到的方法，请告诉我）。

我尝试了以下方法：

>>> np.einsum('ijk,lk->ijk', A, B)
[[[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]

 [[3. 3.]
  [3. 3.]
  [3. 3.]]]

如您所见，这显然与上面的输出不同。

我能想到的另一个解决方案是：

>>> A*B[:,np.newaxis,:].repeat(3, axis=1)
[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

它给出了正确的输出，但我仍然很想知道如何使用 np.einsum

编辑： Warren Weckesser 在评论中指出，上面的解决方案可以简化为 A*B[:,np.newaxis,:]，这是我见过的最干净的解决方案，没有使用 np.einsum.

Answer 1

numpy.einsum解法：

C = np.einsum('ijk,jk->jik', A, B)

使用省略号：

C = np.einsum('ij...,j...->ji...', A, B)

输出

[[[0. 0.]
  [0. 0.]
  [0. 0.]]

 [[1. 1.]
  [1. 1.]
  [1. 1.]]

 [[2. 2.]
  [2. 2.]
  [2. 2.]]]

通过使用 np.einsum 将 2D 数组的行广播到 3D 数组的每个切片进行乘法

Multiplication by broadcasting rows of 2D array to each slice of 3D array using np.einsum

python

arrays

numpy

broadcast

multiplication