将 numpy 3d 数组中的第 i 个 2-d 矩阵与 2d 数组中的第 i 列相乘答案

【问题标题】：Multiply i-th 2-d matrix in numpy 3d array with i-th column in 2d array将 numpy 3d 数组中的第 i 个 2-d 矩阵与 2d 数组中的第 i 列相乘
【发布时间】：2020-11-29 09:14:46
【问题描述】：

假设我有一个 3d 数组 A 和一个 2d 数组 B。A 的维度为 (s,m,m)，而 B 的维度为 (m,s)。

我想为具有维度 (m,s) 的二维数组 C 编写代码，使得 C[:,i] = A[i,:,:] @ B[:,i]。

有没有一种方法可以优雅地做到这一点，而无需在 numpy 中使用 for 循环？

我想到的一个解决方案是将 B 重塑为维度为 (m,s,1) 的 3d 数组，通过 A@B 将 A 和 B 相乘，然后将生成的 3d 数组重塑为 2d 数组。这听起来有点乏味，想知道是否可以在这里应用tensordot或einsum。

建议表示赞赏。谢谢！

【问题讨论】：

标签： python numpy matrix

【解决方案1】：

在这里使用einsum 很简单：

A = np.arange(18).reshape(2,3,3)
B = np.arange(6).reshape(3,2)

C = np.einsum("ijk,ki->ji",A,B)

for i in range(2):
   A[i]@B[:,i]==C[:,i]

# array([ True,  True,  True])
# array([ True,  True,  True])

【讨论】：

非常感谢！这看起来非常干净和高效。