Search code examples
pythonnumpymultidimensional-arraymatrix-multiplicationcupy

Multiply 2D array by each element from 1D array to obtain a 3D array without usage of loops


We have two NumPy arrays with different shapes (n,n) and (m,):

A = [[1 2 3],
 [4 5 6],
 [7 8 9]]
B = [1 2 3 4]

I would like to multiply the 2D array A by each element from the 1D array B to obtain a new 3D matrix like:

C = [
[[1*1 2*1 3*1],
[4*1 5*1 6*1],
[7*1 8*1 9*1]],

[[1*2 2*2 3*2],
[4*2 5*2 6*2],
[7*2 8*2 9*2]],

[[1*3 2*3 3*3],
[4*3 5*3 6*3],
[7*3 8*3 9*3]],

[[1*4 2*4 3*4],
[4*4 5*4 6*4],
[7*4 8*4 9*4]]]

Is it possbile to perform this type of multiplication using NumPy?

I have tried different methods using numpy.reshape(), however I couldn't manage to get the expected result

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


Solution

  • You can also use np.multiply.outer:

    >>> A = np.arange(1, 10).reshape(3, 3)
    >>> B = np.arange(1, 5)
    
    >>> np.multiply.outer(B, A)
    array([[[ 1,  2,  3],
            [ 4,  5,  6],
            [ 7,  8,  9]],
    
           [[ 2,  4,  6],
            [ 8, 10, 12],
            [14, 16, 18]],
    
           [[ 3,  6,  9],
            [12, 15, 18],
            [21, 24, 27]],
    
           [[ 4,  8, 12],
            [16, 20, 24],
            [28, 32, 36]]])