python - 在 numpy 中，将 3 维数组的第二维乘以 1 维数组的最快方法是什么？

Question

你有一个形状为 (a,b,c) 的数组，你想将第二个维度乘以一个形状为 (b) 的数组

for 循环可以，但有更好的方法吗？

例。

A = np.array(shape=(a,b,c))
B = np.array(shape=(b))

for i in B.shape[0]:
    A[:,i,:]=A[:,i,:]*B[i]

Answer 1

使用broadcasting :

A*B[:,np.newaxis]

例如:

In [47]: A=np.arange(24).reshape(2,3,4)

In [48]: B=np.arange(3)

In [49]: A*B[:,np.newaxis]
Out[49]: 
array([[[ 0,  0,  0,  0],
        [ 4,  5,  6,  7],
        [16, 18, 20, 22]],

       [[ 0,  0,  0,  0],
        [16, 17, 18, 19],
        [40, 42, 44, 46]]])

B[:,np.newaxis] 的形状为 (3,1)。广播在左侧添加新轴，所以这是广播形状(1,3,1)。广播还会沿长度为 1 的轴重复项目。因此，当与 A 相乘时，它会进一步广播到形状 (2,3,4)。这与 A 的形状匹配。然后像往常一样按元素进行乘法运算。