使用Pandas计算按钮连续状态的持续时长问题

问题背景 #

原始DataFrame如下：

value  time  button
 10    1000   up
 8     1100   up
 4     1150   down
 1     1200   down
 6     1225   up
 5     1300   no
 5     1350   no

需求：计算按钮处于up/down/no连续状态的持续时长，当状态变化时，用该状态最后一条记录的time减去第一条记录的time，期望结果：

button  time
up       100
down     100
up       25
no       125

错误代码分析 #

你尝试的代码逻辑有误，导致结果全为0：

df['same_as_prev'] = (df['button'] == df['button'].shift()).astype(int)
df['consec_rows'] = df['same_as_prev'].cumsum()
df_result = df.groupby(['button','consec_rows']).agg({'time': lambda x: x.iloc[-1]-x.iloc[0]})

问题出在consec_rows的生成逻辑：当按钮状态不变时，same_as_prev为1，累加后会让同一个连续状态的行分到不同的consec_rows组中（比如连续两行up会被分到0和1组），最终每个分组只有一行数据，时间差自然为0。

正确解决方案 #

核心是正确标记连续状态的分组：当按钮状态与前一行不同时，生成新的分组ID，然后按分组+按钮状态计算时间差。

完整代码：

import pandas as pd

# 构造原始DataFrame
data = {
    'value': [10, 8, 4, 1, 6, 5, 5],
    'time': [1000, 1100, 1150, 1200, 1225, 1300, 1350],
    'button': ['up', 'up', 'down', 'down', 'up', 'no', 'no']
}
df = pd.DataFrame(data)

# 生成连续状态的分组ID：状态变化时分组号+1
df['group'] = (df['button'] != df['button'].shift()).cumsum()

# 按分组和按钮状态计算持续时长
df_result = df.groupby(['group', 'button'])['time'].agg(lambda x: x.iloc[-1] - x.iloc[0]).reset_index()

# 调整结果格式，去掉group列
df_result = df_result[['button', 'time']].reset_index(drop=True)

print(df_result)

运行后得到的结果与期望一致：

button  time
0     up   100
1   down   100
2     up    25
3     no   125

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