Multiplying 2d array by 1d array

2020-04-15 11:29发布

问题:

I have an 2D-array a with shape (k,n) and I want to 'multiply' it with an 1D-array b of shape (m,):

a = np.array([[2, 8],
              [4, 7],
              [1, 2],
              [5, 2],
              [7, 4]])

b = np.array([3, 5, 5])

As a result of the 'multiplication' I'm looking for:

array([[[2*3,2*5,2*5],[8*3,8*5,8*5]],
       [[4*3,4*5,4*5],[7*3,7*5,7*5]],
       [[1*3,1*5,1*5], ..... ]],
          ................. ]]])

= array([[[ 6, 10, 10],
          [24, 40, 40]],

         [[12, 20, 20],
          [21, 35, 35]],

         [[ 3,  5,  5],
          [ ........ ]],

             ....... ]]])

I could solve it with a loop of course, but I'm looking for a fast vectorized way of doing it.

回答1:

Extend a to a 3D array case by adding a new axis at the end with np.newaxis/None and then do elementwise multiplication with b, bringing in broadcasting for a vectorized solution, like so -

b*a[...,None]

Sample run -

In [19]: a
Out[19]: 
array([[2, 8],
       [4, 7],
       [1, 2],
       [5, 2],
       [7, 4]])

In [20]: b
Out[20]: array([3, 5, 5])

In [21]: b*a[...,None]
Out[21]: 
array([[[ 6, 10, 10],
        [24, 40, 40]],

       [[12, 20, 20],
        [21, 35, 35]],

       [[ 3,  5,  5],
        [ 6, 10, 10]],

       [[15, 25, 25],
        [ 6, 10, 10]],

       [[21, 35, 35],
        [12, 20, 20]]])