Использование numpy.where () для создания нового массива

Question

Скажем, у меня есть 2 массива A, B:

A имеет форму (2, 400), что составляет 400 (x, y) точек

B имеет форму (1, 400) которые равны 0 или 1 для каждой из 400 (x, y) точек

Я хочу создать новый массив C такой, что:

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

Итак, форма C должно быть (400, 3).

Я понимаю, что эффективным способом было бы использовать np.where(), но я не могу понять, как это сделать для двумерной матрицы.

Answer 1

Вам не нужен np.where, вместо этого используйте multiply или matmul, так как у вас уже есть маска.Я сделал трюк с: np.matmul(B,np.ones((1,3))*10)+10 или np.multiply(B,np.ones(3) * 10)+10