pythonnumpyvectorization

How do I vectorize this expression with Numpy?

I am trying to compute the array "matrix" below. How can I do this using vectorized functions of NumPy?

x = np.array([
    [2, 1, 0],
    [1, 1, 0],
    [3, 2, 1],
    [1, 0, 0],
    [2, 3, 0]
])

y = np.array([
    [405, 200, 150],
    [200, 300, 150],
    [150, 100, 105],
    [425, 200, 250],
    [500, 620, 300]
])

matrix = np.zeros((5,3,4))  
for i in range(5):
    for j in range(3):
       matrix[i,j,x[i,j]] = y[i,j]

I tried this:

vmatrix = np.zeros((5,3,4))

vmatrix[:,:,x] = y

But it does not work...

Solution

  • You can do this in one line by turning x into one-hot vectors (using methods inspired by this answer) and reshaping y to make the dimensions ready for broadcasting:

    matrix2 = np.eye(4)[x] * np.expand_dims(y, axis=-1)

print(np.equal(matrix, matrix2).all()) # True