Numpy multiplying one dimension of an N-dimensional matrix

user_7571907 0
Viewed 94

I have a 3D numpy array a with shape [M, N, 3], and a 2D array b with shape [M, N]. I want to multiply each of the 3 values in the 3rd dimension of a by the corresponding value in b, which I am currently doing like so:

a[:,:,0] *= b
a[:,:,1] *= b
a[:,:,2] *= b

I was wondering if there is a single-line way to achieve this.

Thanks!

numpy python
1 Answers

a * b[:,:,np.newaxis] would do it:

In [37]: a
Out[37]:
array([[[1, 2, 3],
        [4, 5, 6]]])

In [38]: b
Out[38]: array([[4, 5]])

In [39]: a * b[:,:,np.newaxis]
Out[39]:
array([[[ 4,  8, 12],
        [20, 25, 30]]])

There are other ways to express the same thing, such as

  • a * np.expand_dims(b, 2).
  • a * b[..., np.newaxis]
user_3672730
Related