开发者

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

开发者 https://www.devze.com 2023-03-29 06:41 出处:网络
You have an array of shape (a,b,c) and you want to multiply the second dimension by an array of shape (b)

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.

0

精彩评论

暂无评论...
验证码 换一张
取 消

关注公众号