In numpy, what is the fastest way to multiply the second dimension of a 3 dimensional array by a 1 dimensional array?

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

A for loop would 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]

Use broadcasting:

A*B[:,np.newaxis]

For example:

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 shape (3,1). Broadcasting adds new axes on the left, so this is broadcasted to shape (1,3,1). Broadcasting also repeats the items along axes with length 1. So when multiplied with A, it gets further broadcasted to shape (2,3,4). This matches the shape of A. Multiplication then proceeds element-wise, as always.