如何基于DataFrame分组，通过Start与End值计算Density列？

原始数据 #

Rate     Distance    
Start           4                   
Coupon          7                   
Coupon          8                   
End             10                 
Start           13                  
Coupon          14                  
End             18                 

需求说明 #

新增Density列，计算规则：

• 以Start到End为一个分组区间
• 每组内：Density = 组内Coupon数量 / (End的Distance值 - Start的Distance值)
• 将计算结果填充到该分组的所有行中

实现代码（Pandas） #

import pandas as pd

# 构造原始DataFrame
df = pd.DataFrame({
    "Rate": ["Start", "Coupon", "Coupon", "End", "Start", "Coupon", "End"],
    "Distance": [4, 7, 8, 10, 13, 14, 18]
})

# 生成分组标识：每出现一次Start，分组ID加1
df["group_id"] = df["Rate"].eq("Start").cumsum()

# 分组计算每组的Coupon数量、Start与End的Distance差值
group_metrics = df.groupby("group_id").agg(
    coupon_num=("Rate", lambda x: x[x == "Coupon"].count()),
    dist_diff=("Distance", lambda x: x[df.loc[x.index, "Rate"] == "End"].iloc[0] - x[df.loc[x.index, "Rate"] == "Start"].iloc[0])
)

# 计算每组的Density并映射回原表
df["Density"] = df["group_id"].map(group_metrics.apply(lambda row: row["coupon_num"] / row["dist_diff"], axis=1))

# 整理输出格式，保留目标列并保留两位小数
result_df = df[["Rate", "Distance", "Density"]].round(2)
print(result_df.to_string(index=False))

输出结果 #

Rate  Distance  Density
Start         4     0.33
Coupon        7     0.33
Coupon        8     0.33
End          10     0.33
Start        13     5.00
Coupon       14     5.00
End          18     5.00

关键步骤说明 #

1. 分组标识生成：通过eq("Start").cumsum()实现，每遇到一个Start就创建新的分组ID，自动划分Start-End区间。
2. 分组统计：对每个分组统计Coupon的数量，同时提取组内Start和End对应的Distance值计算差值。
3. 密度计算与映射：计算每组的密度值后，通过分组ID将结果映射到组内所有行，保证同一分组的Density值一致。

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