как установить значение в 2 ячейке, если другая ячейка содержит «что-то»

Question

У меня есть несколько панд данных:

a,b,c
AAA,,
DDD,,
KKK,,
AAA,,  

Я хочу выполнить поиск в столбце «A», и если строка в столбце «A» содержит слово «AAA», мне нужно установить значение «BBB» в столбце «B» и «CCC» в столбце «C».
Итак, я хочу получить результат следующим образом:

a,b,c
AAA,BBB,CCC
DDD,,
KKK,,
AAA,BBB,CCC

Я написал код с numpy:

df['b'] = pd.np.where(df.a.str.contains("AAA"), "BBB", '')

Как расширить его для работы со столбцами 'b' и 'c'?

Answer 1

Вы можете использовать двойной np.where:

mask = df.a.str.contains("AAA")
df['b'] = pd.np.where(mask, "BBB", '')
df['c'] = pd.np.where(mask, "CCC", '')

Или assign:

mask = df.a.str.contains("AAA")
df = df.assign(b=pd.np.where(mask, "BBB", ''), c=pd.np.where(mask, "CCC", ''))

Если необходимо создать несколько столбцов с одним np.where, необходимо создать Nx1 mask:

mask = df.a.str.contains("AAA")[:, None]
df[['b','c']] = np.where(mask, ['BBB','CCC'], ['',''])
print (df)
     a    b    c
0  AAA  BBB  CCC
1  DDD          
2  KKK          
3  AAA  BBB  CCC