Python:pandas groupby ID，添加多列的值，直到满足中断条件

Question

我有一个数据框

df = pd.DataFrame({ 'id': ['1','1','1','2','2','3','4'],
'transfer_time': [0,10,12,0,19,0,0],
'trip_time': [20,30,9,10,15,12,16],
'is_final_stop': [0,1,0,0,1,1,1]
})
print(df)

如下所示:

id  is_final_stop   transfer_time   trip_time
1   0                 0               20
1   1                 10              30
1   0                 12              9
2   0                 0               10
2   1                 19              15
3   1                 0               12
4   1                 0               16

我需要计算旅程时间 = 转移时间 + 行程时间，按 id 分组，直到 is_final_stop 标志 = 1。(不应考虑将此实例之后的 id 添加到旅程时间)

解决方案应如下所示:

id journey_time
1  60
2  44
3  12
4  16

希望有一个能够在 400 万行数据帧上高效工作的解决方案。干杯。

Answer 1

通过反向cumsum计算“包含”标志，然后在聚合时间上使用GroupBy + sum :

inc_flag = df.iloc[::-1].groupby('id')['is_final_stop'].cumsum().iloc[::-1].astype(bool)

res = df[inc_flag].groupby('id')[['transfer_time', 'trip_time']].sum().sum(1)\
                  .rename('total_time').reset_index()

结果:

  id  total_time
0  1          60
1  2          44
2  3          12
3  4          16