matrix multiplication in a customized way

I have a array of the size (3,):

a = [1,2,3]

and another array of the size (3,3):

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

I am looking for a vectorized way to multiply the a**2 (a^2) to b in a way to:

> a**2=[1,4,9]

multiplies that 1 to entire row 1 from matrix b
4 to the second row of the matrix b
and 9 to the entire row of the matrix b.

my final result has to be this:

> (a**2)*b = [[1,2,3],[16,20,24],[63,72,81]]

Thanks!

>Solution :

First, you need to convert your arrays into numpy arrays:

import numpy as np
a = [1,2,3]
b = [[1,2,3],[4,5,6],[7,8,9]]
a, b = np.array(a), np.array(b)

You can use:

(a**2)[:, None] * b

Output:

array([[ 1,  2,  3],
       [16, 20, 24],
       [63, 72, 81]])