python - 如何使用 groupby 减去列中的值

Question

我有以下数据框:

ID  Days TreatmentGiven TreatmentNumber
--- ---- -------------- ---------------
1    0      False             NaN
1    30     False             NaN
1    40     True               1
1    56     False             NaN 
2    0      False             NaN
2    14     True               1
2    28     True               2 

我想根据第一次治疗的时间 (TreatmentNumber==1) 创建一个新列，其中包含天数的新基线，并按 ID 分组，以便结果如下:

ID  Days TreatmentGiven TreatmentNumber New_Baseline
--- ---- -------------- --------------- ------------
1    0      False             NaN          -40
1    30     False             NaN          -10
1    40     True               1            0
1    56     False             NaN           16
2    0      False             NaN          -14
2    14     True               1            0
2    28     True               2            14

执行此操作的最佳方法是什么？

谢谢。

Answer 1

这是一种使用 series.where + groupby+transform 的方法:

s = df['Days'].where(df['TreatmentGiven']).groupby(df['ID']).transform('first')
df['New_Baseline'] = df['Days'].sub(s)

---

输出

   ID  Days  TreatmentGiven  TreatmentNumber  New_Baseline
0   1     0           False              NaN         -40.0
1   1    30           False              NaN         -10.0
2   1    40            True              1.0           0.0
3   1    56           False              NaN          16.0
4   2     0           False              NaN         -14.0
5   2    14            True              1.0           0.0
6   2    28            True              2.0          14.0