使用np.einsum将二维阵列的行广播到三维阵列的每一个片断，进行乘法。

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