Warm tip: This article is reproduced from serverfault.com, please click

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

(python - Multiply i-th 2-d matrix in numpy 3d array with i-th column in 2d array)

发布于 2020-11-29 00:46:11

假设我有一个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。

建议表示赞赏。谢谢!

Questioner
secondrate
Viewed
0
Paul Panzer 2020-11-29 09:14:46

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