如何快速计算DataFrame中各分组缺失的累计和(end_time)

需求说明 #

给定包含idx分组列的DataFrame（同idx行相邻），需高效计算每组最后一行的end_time值：

• 除最后一行外的end_time可通过df["end_time"] = df["start_time"].shift(periods=-1).fillna(0).astype(float)得到
• 每组最后一行的end_time为该组dt的累计和加上组内首个start_time（等价于组内倒数第二行end_time加最后一行dt）

高效实现方法（无需循环） #

方法1：直接生成所有行的end_time（推荐，全向量操作）

这种方法无需先计算shift结果，一步生成所有行的end_time，效率最高：

import pandas as pd

# 示例DataFrame
data = {
    "idx": [1,1,1,2,2,2],
    "start_time": [0,1,3,0,2,6],
    "dt": [1,2,3,2,4,5]
}
df = pd.DataFrame(data)

# 计算组内dt的累计和
df["cum_dt"] = df.groupby("idx")["dt"].cumsum()
# 获取每组的首个start_time
df["first_start"] = df.groupby("idx")["start_time"].transform("first")
# 计算所有行的end_time
df["end_time"] = df["first_start"] + df["cum_dt"]

# 可选择删除中间列
df.drop(["cum_dt", "first_start"], axis=1, inplace=True)

运行后得到的结果：

idx  start_time  dt  end_time
0    1           0   1       1.0
1    1           1   2       3.0
2    1           3   3       6.0
3    2           0   2       2.0
4    2           2   4       6.0
5    2           6   5      11.0

方法2：基于已有shift结果补全最后一行

如果已经通过shift生成了除最后一行外的end_time，可以用以下方式快速补全：

# 先执行你提供的shift操作生成部分end_time
df["end_time"] = df["start_time"].shift(periods=-1).fillna(0).astype(float)

# 计算每组的总dt和与首个start_time
total_dt = df.groupby("idx")["dt"].transform("sum")
first_start = df.groupby("idx")["start_time"].transform("first")

# 标记每组最后一行并赋值
is_last_row = df.groupby("idx").cumcount(ascending=False) == 0
df.loc[is_last_row, "end_time"] = first_start[is_last_row] + total_dt[is_last_row]

关键说明 #

两种方法均使用Pandas的分组向量操作，避免了逐行循环的低效问题，处理大数据集时优势明显。

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