Multiply each row of one array with each element of another array in numpy

后端 未结 1 1160
甜味超标
甜味超标 2021-01-12 23:01

I have two arrays A and B in numpy. A holds cartesian coordinates, each row is one point in 3D space and has the shape (r, 3). B has the sh

1条回答
  •  余生分开走
    2021-01-12 23:13

    Use broadcasting -

    A[:,None,:]*B[:,:,None]
    

    Since np.einsum also supports broadcasting, you can use that as well (thanks to @ajcr for suggesting this concise version) -

    np.einsum('ij,ik->ikj',A,B)
    

    Sample run -

    In [22]: A
    Out[22]: 
    array([[1, 1, 1],
           [2, 2, 2],
           [3, 3, 3]])
    
    In [23]: B
    Out[23]: 
    array([[10, 20],
           [30, 40],
           [50, 60]])
    
    In [24]: A[:,None,:]*B[:,:,None]
    Out[24]: 
    array([[[ 10,  10,  10],
            [ 20,  20,  20]],
    
           [[ 60,  60,  60],
            [ 80,  80,  80]],
    
           [[150, 150, 150],
            [180, 180, 180]]])
    
    In [25]: np.einsum('ijk,ij->ijk',A[:,None,:],B)
    Out[25]: 
    array([[[ 10,  10,  10],
            [ 20,  20,  20]],
    
           [[ 60,  60,  60],
            [ 80,  80,  80]],
    
           [[150, 150, 150],
            [180, 180, 180]]])
    

    0 讨论(0)
提交回复
热议问题