如何在含diff函数的Pandas groupby代码中添加心率平均值列？

原代码与需求 #

原代码用于计算高功率（功率>200）区间的时长，现需新增区间内的平均功率、平均心率列，期望输出包含这三个指标的DataFrame。

原代码：

import pandas as pd
from pandas import Timestamp

data= pd.DataFrame(
    [[Timestamp('2022-08-05 10:11:04'), 140, 120],
    [Timestamp('2022-08-05 10:11:05'), 160, 155],
    [Timestamp('2022-08-05 10:11:06'), 230, 156],
    [Timestamp('2022-08-05 10:11:07'), 230, 155],
    [Timestamp('2022-08-05 10:11:08'), 230, 160],
    [Timestamp('2022-08-05 10:11:09'), 140, 130],
    [Timestamp('2022-08-05 10:11:10'), 140, 131],
    [Timestamp('2022-08-05 10:11:11'), 230, 170]],
    columns=['timestamp', 'power', 'heart rate'])

m = data['power'].gt(200) # 标记高功率行
gb = (-data['timestamp'].diff(-1))[m].groupby([(~m).cumsum()]).sum()
gb= gb.groupby((~m).cumsum()).sum()
gb

最佳实现方式 #

直接利用groupby.agg()一次性聚合所有需要的指标，无需拆分计算，代码更简洁高效，步骤如下：

1. 生成分组键：(~m).cumsum()用于将连续的高功率行划分为同一组；
2. 过滤出高功率数据行，按分组键分组；
3. 对每组同时计算时长、平均功率、平均心率；
4. 重命名列名以匹配期望输出。

优化后代码：

import pandas as pd
from pandas import Timestamp

data= pd.DataFrame(
    [[Timestamp('2022-08-05 10:11:04'), 140, 120],
    [Timestamp('2022-08-05 10:11:05'), 160, 155],
    [Timestamp('2022-08-05 10:11:06'), 230, 156],
    [Timestamp('2022-08-05 10:11:07'), 230, 155],
    [Timestamp('2022-08-05 10:11:08'), 230, 160],
    [Timestamp('2022-08-05 10:11:09'), 140, 130],
    [Timestamp('2022-08-05 10:11:10'), 140, 131],
    [Timestamp('2022-08-05 10:11:11'), 230, 170]],
    columns=['timestamp', 'power', 'heart rate'])

m = data['power'].gt(200)
# 生成分组键，区分不同的高功率区间
group_key = (~m).cumsum()

# 过滤高功率数据，分组后聚合多指标
gb = data[m].groupby(group_key).agg(
    time_at_high_intensity=('timestamp', lambda x: (x.max() - x.min())),
    avg_power=('power', 'mean'),
    avg_heart_rate=('heart rate', 'mean')
).reset_index(drop=True)

# 如果需要将时长转换为HH:MM:SS格式（如期望输出），可添加以下代码
gb['time_at_high_intensity'] = pd.to_datetime(gb['time_at_high_intensity'].dt.total_seconds(), unit='s').dt.time

print(gb)

代码说明 #

• 原代码中重复的groupby操作可直接省略，一次groupby.agg()即可完成所有计算；
• 时长计算用x.max() - x.min()替代原diff(-1)求和的方式，逻辑更直观，结果一致；
• 若期望时长以HH:MM:SS格式显示，可通过pd.to_datetime将秒数转换为时间格式；
• 聚合时直接指定列名和对应的计算逻辑，一步得到目标DataFrame。

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