python - 我如何在没有循环的情况下在 NumPy 中实现这个 "row in A * all rows in B by col in A"乘法？

Question

假设我有两个矩阵 A 和 B:

A = np.array([[1, 3, 2],
              [2, 2, 3],
              [3, 1, 1]])

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

我想取 A 中的一列并将其乘以 B 中的每一列元素，然后继续 A 中的下一列。因此，仅以一列为例，我将使用 A[:,0] (值 1,2,3)，然后乘以 B 中的每一列得到:

array([[0, 1, 0],
       [2, 2, 0],
       [3, 3, 3]])

我已经使用 np.einsum 实现了这个像这样:

np.einsum('i,ij->ij',A[:,0],B)

然后我想生成一个 3D 矩阵，其深度维度对应于 A 中每一列的乘法，我使用 for 循环实现:

np.stack([np.einsum('i,ij->ij',A[:,i],B) for i in range(0,A.shape[1])])

这将返回我想要的数组:

array([[[0, 1, 0],
        [2, 2, 0],
        [3, 3, 3]],

       [[0, 3, 0],
        [2, 2, 0],
        [1, 1, 1]],

       [[0, 2, 0],
        [3, 3, 0],
        [1, 1, 1]]])

如果没有循环，我将如何进行此操作？这可以完全通过 np.einsum 完成吗？ ? NumPy 中是否有另一个函数可以更简单地做到这一点？

Answer 1

这是一个简单的方法:

A.T[:,:,None]*B

添加最后一个 None在索引中创建一个新轴，然后用于广播元素乘法。