将2d数组乘以1d数组

时间:2016-01-11 08:02:14

标签: python arrays performance numpy vectorization

我有一个带有形状(k,n)的二维数组a,我想将它与一个形状为(m,)的一维数组b'相乘':

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

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

由于'乘法'我正在寻找:

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],
          [ ........ ]],

             ....... ]]])

我当然可以用循环来解决它,但我正在寻找一种快速的矢量化方法。

1 个答案:

答案 0 :(得分:1)

通过在np.newaxis/None末尾添加一个新轴,然后使用a进行元素乘法,将broadcasting扩展为一个3D数组大小写,为b引入{{3}}解决方案,像这样 -

b*a[...,None]

示例运行 -

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]]])