групповые операции с условиями в панде

Question

Я хочу выполнить групповую операцию в пандах. Например, я хочу сгруппировать столбец patient, а если столбец treatment == X перенести соответствующее значение doctor в новый столбец с именем nurse.

Например: df

import pandas as pd
import numpy as np

df = pd.DataFrame({'patient': ['a','a','a','b','b','b'],
   ....:           'treatment': ['X','Y','Y','X','Z','Z'],
                   'doctor': ['1','2','2','2','3','3']})

  patient treatment doctor
0       a         X      1
1       a         Y      2
2       a         Y      2
3       b         X      2
4       b         Z      3
5       b         Z      3

Я пытался

df=df.assign(nurse=np.where(df.['treatment'].str.contains('X'),df.groupby('patient')['doctor'], np.nan))

, но получал ошибку

SyntaxError: недопустимый синтаксис

ожидаемый результат

    patient treatment doctor  nurse
0       a         X      1      1
1       a         Y      2      1
2       a         Y      2      1
3       b         X      2      2
4       b         Z      3      2
5       b         Z      3      2

Как мне добиться этого результата?

thx

Answer 1

Использование DataFrame.apply + Series.where . Затем фаршированные ffill :

df['nurse']=df.groupby('patient',sort=False).apply(lambda x: x['doctor'].where(x['treatment'].eq('X')).ffill()).reset_index(drop=True)
print(df)

---

     patient treatment doctor nurse
0       a         X      1     1
1       a         Y      2     1
2       a         Y      2     1
3       b         X      2     2
4       b         Z      3     2
5       b         Z      3     2