In numpy, what is the fastest way to multiply the second dimension of a three-dimensional array by a one-dimensional array?

Question

You have an array of shapes (a, b, c) and you want to multiply the second dimension by an array of shapes (b)

And for the loop it will work, but is there a better way?

Ex.

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

Use broadcasting :

 A*B[:,np.newaxis] 

For instance:

 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] has the form (3.1). Broadcasting adds new axes on the left, so it is broadcast in the form of (1,3,1). Broadcasting also repeats elements along axes of length 1. Therefore, when it is multiplied by A , it is then translated into form (2,3,4). This corresponds to form A Then multiplication continues by elements, as always.