Python优化基于日历年的滚动5年平均值计算方案问询

需求说明 #

现有包含多series标识的周度数据DataFrame，每个series每周均有记录，但每年对应日期不固定。需按series分组，为每条记录计算滚动5年平均值：取该记录过去5年中，日历年（dayofyear）与当前记录最接近的对应值参与平均。

示例数据 #

import pandas as pd
import numpy as np

np.random.seed(42)
dates = pd.date_range(start='2018-01-01', end='2023-12-31', freq='W')
series_ids = ['A', 'B', 'C']

df = pd.DataFrame({
    'series': np.random.choice(series_ids, size=len(dates)),
    'period': dates,
    'value': np.random.randn(len(dates))
}).sort_values(['series', 'period']).reset_index(drop=True)

当前实现（效率较低版本） #

def rolling_5y_avg(row, group_df):
    current_doy = row['period'].dayofyear
    current_year = row['period'].year
    past_5y = group_df[group_df['period'].dt.year.between(current_year-5, current_year-1)]
    if past_5y.empty:
        return np.nan
    
    yearly_matches = []
    for year in past_5y['period'].dt.year.unique():
        year_data = past_5y[past_5y['period'].dt.year == year]
        year_data['doy_diff'] = abs(year_data['period'].dt.dayofyear - current_doy)
        closest = year_data.loc[year_data['doy_diff'].idxmin(), 'value']
        yearly_matches.append(closest)
    return np.mean(yearly_matches)

df['rolling_5y_avg'] = df.groupby('series').apply(
    lambda g: g.apply(rolling_5y_avg, group_df=g, axis=1)
).reset_index(level=0, drop=True)

优化后的高效实现 #

利用pandas的向量化分组与合并操作，替代逐行循环，大幅提升大样本下的运行效率：

# 提取日历年和年份字段
df['doy'] = df['period'].dt.dayofyear
df['year'] = df['period'].dt.year

# 预处理：按series+year存储所有记录的doy和value，便于后续匹配
yearly_data = df.groupby(['series', 'year'])[['doy', 'value']].apply(lambda x: x.reset_index(drop=True))

def compute_rolling_avg(group):
    # 为每条记录生成过去5年的待匹配年份范围
    group['target_years'] = group['year'].apply(lambda y: list(range(y-5, y)))
    # 展开为每条记录对应多个待匹配年份的行
    exploded = group.explode('target_years').rename(columns={'year': 'current_year', 'target_years': 'year'})
    # 合并预处理的年度数据
    merged = exploded.merge(yearly_data, on=['series', 'year'], suffixes=('_current', '_match'))
    # 计算日历年差值，筛选每个原记录对应的最接近匹配项
    merged['doy_diff'] = abs(merged['doy_current'] - merged['doy_match'])
    closest_matches = merged.loc[merged.groupby(['series', 'period'])['doy_diff'].idxmin()]
    # 分组计算平均值并关联回原数据
    rolling_avg = closest_matches.groupby(['series', 'period'])['value_match'].mean().rename('rolling_5y_avg')
    return group.join(rolling_avg)

# 分组执行并整理结果
df_optimized = df.groupby('series').apply(compute_rolling_avg).reset_index(drop=True)
# 清理临时字段
df_optimized = df_optimized.drop(['doy', 'year', 'target_years'], axis=1)

优化核心点 #

1. 避免逐行循环：用explode+merge替代自定义行级循环，利用pandas内置的向量化运算提升速度
2. 预处理复用：提前按series+year存储年度数据，避免重复筛选
3. 高效匹配：通过groupby+idxmin快速定位每个记录的最接近日历年匹配项，时间复杂度远低于逐行遍历

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