Выберите строки на основе условия и установите значения из вектора

Question

Я хочу установить целые строки в значение из вектора, если выполняется условие в столбце.

import pandas as pd

df = pd.DataFrame([['a', 1, 1], ['a', 1, 1], ['b', 1, 1]], columns=('one', 'two', 'three'))
vector = pd.Series([2,3,4])
print(df)

  one  two  three
0   a    1      1
1   a    1      1
2   b    1      1

Я хочу, чтобы результат был таким:

df_wanted = pd.DataFrame([['a', 1, 1], ['a', 1, 1], ['b', 4, 4]], columns=('one', 'two', 'three'))
print(df_wanted)

  one  two  three
0   a    1      1
1   a    1      1
2   b    4      4

Я пробовал это, но выдает ошибку:

df.loc[df['one']=='b'] = vector[df['one']=='b']
ValueError: Must have equal len keys and value when setting with an iterable

// m.

Answer 1

Вы можете указать столбцы в списке для набора:

df.loc[df['one']=='b', ['two', 'three']] = vector[df['one']=='b']
print(df)
  one  two  three
0   a    1      1
1   a    1      1
2   b    4      4

Или, если нужно более динамичное решение - выберите все числовые столбцы:

df.loc[df['one']=='b', df.select_dtypes(np.number).columns] = vector[df['one']=='b']

Или сравните только один раз и присвойте переменной:

m = df['one']=='b'
df.loc[m, df.select_dtypes(np.number).columns] = vector[m]