如何计算同阶段连续分组时长并为每行填充对应组时长？

原始数据 #

dt_datetime        stage    proc_val
2011-11-13 11:00   0        20
2011-11-13 11:10   0        21
2011-11-13 11:30   1        25
2011-11-13 11:40   2        22
2011-11-13 11:55   2        28
2011-11-13 12:00   2        29

期望结果 #

需要新增stage_duration列，同一连续阶段的所有行都填充该阶段的总持续时长（分钟）：

dt_datetime        stage    proc_val   stage_duration
2011-11-13 11:00   0        20         30
2011-11-13 11:10   0        21         30
2011-11-13 11:30   1        25         10
2011-11-13 11:40   2        22         20
2011-11-13 11:55   2        28         20
2011-11-13 12:00   2        29         20

用户尝试的错误代码 #

以下代码无法得到预期结果，因为它仅计算同stage行的时间差累加，未区分连续的stage分组：

df['stage_duration'] = df.groupby('stage')['dt_datetime'].diff().dt.total_seconds() / 60
df['stage_duration'] = df['stage_duration'].cumsum()

补充要求 #

解决方案需支持非连续同stage场景：同一个stage多次出现但不连续时，各自计算持续时长，例如：

dt_datetime        stage    proc_val   stage_duration
2011-11-13 11:00   0        20         30
2011-11-13 11:10   0        21         30
2011-11-13 11:30   1        25         10
2011-11-13 11:40   2        22         20
2011-11-13 11:55   2        28         20
2011-11-13 12:00   2        29         20
2011-11-13 12:00   0        20         70
2011-11-13 13:10   0        21         70

核心思路 #

先识别连续的同stage分组（而非所有同stage的行），计算每个连续分组的总持续时长后，再将时长填充到该分组的所有行中。

实现代码 #

import pandas as pd

# 确保时间列是datetime类型（若原始数据未转换）
df['dt_datetime'] = pd.to_datetime(df['dt_datetime'])

# 生成连续stage的分组ID：当stage发生变化时，分组ID递增
df['stage_group'] = (df['stage'] != df['stage'].shift()).cumsum()

# 按分组ID计算每个连续阶段的持续时长（单位：分钟）
group_duration = df.groupby('stage_group')['dt_datetime'].agg(
    lambda x: (x.max() - x.min()).total_seconds() / 60
).reset_index(name='stage_duration')

# 将计算好的时长合并回原数据，删除临时分组列
df = df.merge(group_duration, on='stage_group').drop('stage_group', axis=1)

代码说明 #

1. 生成连续分组ID：(df['stage'] != df['stage'].shift()).cumsum() 通过对比当前行与上一行的stage值生成布尔序列，累加后得到唯一的连续分组ID，确保不连续的同stage被分到不同组。
2. 计算分组时长：对每个连续分组，取组内最大时间（阶段结束时间）与最小时间（阶段开始时间）的差值，转换为分钟得到该阶段的总持续时长。
3. 填充时长到原数据：通过merge将每个分组的时长对应填充到组内所有行，最后删除临时的分组ID列即可。

内容的提问来源于stack exchange，提问作者Fluxy