Using numpy.where() to create a new array

Question

Say I have 2 arrays A , B :

A has shape (2, 400) which are 400 (x, y) points

B has shape (1, 400) which are either 0 or 1 for each of the 400 (x, y) points

I want to create a new array C such that:

If B[i] == 0, C[i] = [10, 10, 10]
If B[i] == 1, C[i] = [20, 20, 20]

So the shape of C should be (400, 3).

I understand that an efficient way would be to use np.where() but I cannot figure out how to do this for a 2D matrix.

Answer 1

You don't need np.where, instead, use multiply or matmul since you already have a mask. I did the trick with: np.matmul(B,np.ones((1,3))*10)+10 or np.multiply(B,np.ones(3) * 10)+10