Изменить значения в одном массиве на основе значения из столбца второго массива

Question

У меня есть одномерный массив с именем Y_train, который содержит серии из 1 и 0. У меня есть другой массив с именем sample_weight, который представляет собой массив всех 1, имеющий форму Y_train, определяемый как:

sample_weight = np.ones(Y_train.shape, dtype=int)

Я пытаюсь изменить значения в sample_weight на 2, где соответствующее значение в Y_train == 0. Итак, изначально бок о бок это выглядит так:

Y_train        sample_weight
0              1
0              1
1              1
1              1
0              1
1              1

, и я хотел бы, чтобы после преобразования это выглядело так:

Y_train        sample_weight
0              2
0              2
1              1
1              1
0              2
1              1

Я пытался использовать for l oop (показано ниже), но ни одна из 1 не меняется на 2 в sample_weight. Я хотел бы как-то использовать функцию np.where(), если это возможно, но это не принципиально, просто хотелось бы избежать for l oop:

sample_weight = np.ones(Y_train.shape, dtype=int)
for num, i in enumerate(Y_train):
    if i == 0:
        sample_weight[num] == 2

Я пытался использовать показанное решение здесь но безуспешно со вторым массивом. Любые идеи??? Спасибо!

Answer 1

import numpy as np

Y_train = np.array([0,0,1,1,0,1])
sample_weight = np.where(Y_train == 0, 2, Y_train)

>> print(sample_weight)
[2 2 1 1 2 1]

np.where в основном работает так же, как "IF" в Excel:

np.where(condition, then, else)

Работает и для транспонированных массивов:

Y_train = np.array([[0,0,1,1,0,1]]).T
sample_weight = np.where(Y_train == 0, 2, Y_train)

>> print(sample_weight)
[[2]
 [2]
 [1]
 [1]
 [2]
 [1]]