Pandas DataFrame分组取首行并按Trans组计算日期差问题求解

原DataFrame结构 #

import pandas as pd

data = {
    'glob_order': [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11],
    'trans': ['A', 'A', 'A', 'A', 'B', 'B', 'B', 'C', 'C', 'C', 'C'],
    'chain': [1, 1, 2, 2, 1, 1, 2, 1, 1, 2, 2],
    'date': ['1/08/2023', '2/08/2023', '3/08/2023', '4/08/2023', '5/08/2023', '6/08/2023', '7/08/2023', '8/08/2023', '9/08/2023', '10/08/2023', '11/08/2023']
}

df = pd.DataFrame(data)
df['date'] = pd.to_datetime(df['date'], format='%d/%m/%Y')

需求 #

• 按trans和chain列分组，选取每个组的第一行；
• 创建delta列，计算每个trans组内当前行日期与前一行日期的天数差，每个trans组的第一行delta为NaN。

错误尝试及问题 #

之前的代码直接全局计算日期偏移差，导致跨trans组计算日期差，不符合需求：

(df
.groupby(['trans', 'chain'])
.first()
.assign(
 delta=lambda x: (x['date'] - x['date'].shift(1)).dt.total_seconds() / (60*60*24),
 ).reset_index()
)

正确链式解决方案 #

result = (df
          # 第一步：按trans和chain分组，保留原结构取每组第一行
          .groupby(['trans', 'chain'], as_index=False)
          .first()
          # 第二步：按trans分组，组内计算日期差
          .assign(
              delta=lambda x: x.groupby('trans')['date']
                               .diff()
                               .dt.days
          )
)

print(result)

代码说明 #

• groupby(['trans', 'chain'], as_index=False).first()：分组后不将分组列设为索引，直接得到每个组的第一行数据；
• x.groupby('trans')['date'].diff().dt.days：针对每个trans组单独计算日期差，diff()自动为组内第一行返回NaN，dt.days直接提取天数差，无需手动转换秒数，更简洁高效。

期望输出 #

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