Pandas按A字段分组计算各分组首次出现后6-12个月B列的平均值

实现思路 #

• 第一步：按A列分组，计算每个A取值对应的最早观测时间
• 第二步：将每个A对应的最早时间关联到原DataFrame的每一行
• 第三步：筛选出行索引时间落在「最早时间+6个月」到「最早时间+12个月」区间内的有效记录
• 第四步：按A列分组计算B列的平均值，重命名列后得到最终结果

完整实现代码 #

import pandas as pd

# 构造示例数据
data = {
    'A': ['x', 'x', 'x', 'x', 'y', 'y', 'y', 'y', 'y'],
    'B': [10, 32, 12, 13, 24, 32, 12, 72, 90],
    'created_on': [
        '2018-01-31', 
        '2019-02-25', 
        '2018-02-12', 
        '2019-05-31', 
        '2021-03-12', 
        '2020-04-23', 
        '2016-01-11', 
        '2016-05-02', 
        '2018-12-31',
    ]
}
df = pd.DataFrame(data)
df = df.set_index(pd.to_datetime(df['created_on']))
df.drop(['created_on'], axis=1, inplace=True)

# 核心计算逻辑
# 1. 计算每个A的最早出现时间
df_first = df.reset_index().groupby('A', as_index=False)['created_on'].min().rename(columns={'created_on': 'first_occur_dt'})
# 2. 把最早时间回拼到原表
df = df.merge(df_first, on='A', how='left').set_index(df.index)
# 3. 筛选符合时间区间的行
mask = (df.index >= df['first_occur_dt'] + pd.DateOffset(months=6)) & (df.index <= df['first_occur_dt'] + pd.DateOffset(months=12))
df_filter = df[mask]
# 4. 分组求均值，如需要保留无有效数据的A取值，可改为和df_first[['A']]左连接
result = df_filter.groupby('A', as_index=False)['B'].mean().rename(columns={'B': 'avg_B_6_12_months_after_earliest_observation'})

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