如何高效生成基于当前行日期前数据的均值列并优化Pandas代码

原代码使用iterrows()逐行迭代，每次循环都多次对整个DataFrame进行切片筛选，时间复杂度为O(n²)，当数据量达到5万行时，会出现严重的性能瓶颈。以下是优化后的代码，兼顾性能与可读性：

优化后的代码 #

import pandas as pd

# 构造示例数据
data = {
    'id': [1,2,3,4,5],
    'home_goals': [1,0,3,1,2],
    'away_goals': [1,1,2,0,1],
    'home_name': ['a','b','a','b','a'],
    'away_name': ['b','a','b','a','b'],
    'date': ['2020-01-01', '2020-01-02', '2020-01-03', '2020-01-04', '2020-01-05'],
}

df = pd.DataFrame(data=data)

# 1. 转换日期类型并按日期排序（确保历史数据顺序正确）
df['date'] = pd.to_datetime(df['date'])
df = df.sort_values('date').reset_index(drop=True)

# 2. 批量计算分组累积统计量，减少重复分组操作
# 计算球队作为主场的累积进球和场次
home_group = df.groupby('home_name').expanding()
home_group_sum = home_group['home_goals'].sum().reset_index(level=0, drop=True)
home_group_count = home_group.size().reset_index(level=0, drop=True)

# 计算球队作为客场的累积进球和场次
away_group = df.groupby('away_name').expanding()
away_group_sum = away_group['away_goals'].sum().reset_index(level=0, drop=True)
away_group_count = away_group.size().reset_index(level=0, drop=True)

# 3. 构造各统计项的历史数据（shift(1)排除当前行）
# 主队在主场的历史进球统计
home_cum = pd.DataFrame({
    'home_scored_at_home_sum': home_group_sum,
    'home_scored_at_home_count': home_group_count
}).shift(1)

# 客队在客场的历史进球统计
away_cum = pd.DataFrame({
    'away_scored_at_away_sum': away_group_sum,
    'away_scored_at_away_count': away_group_count
}).shift(1)

# 主队作为客场的历史进球统计
home_as_away_cum = pd.DataFrame({
    'home_scored_at_away_sum': away_group_sum,
    'home_scored_at_away_count': away_group_count
}).shift(1)

# 客队作为主场的历史进球统计
away_as_home_cum = pd.DataFrame({
    'away_scored_at_home_sum': home_group_sum,
    'away_scored_at_home_count': home_group_count
}).shift(1)

# 4. 合并所有统计数据到原DataFrame
df = pd.concat([df, home_cum, away_cum, home_as_away_cum, away_as_home_cum], axis=1)

# 5. 计算均值，处理场次为0的情况
df['home_scored_at_home_mean'] = df.apply(
    lambda x: x['home_scored_at_home_sum'] / x['home_scored_at_home_count'] 
    if x['home_scored_at_home_count'] > 0 else None, axis=1
)
df['away_scored_at_away_mean'] = df.apply(
    lambda x: x['away_scored_at_away_sum'] / x['away_scored_at_away_count'] 
    if x['away_scored_at_away_count'] > 0 else None, axis=1
)
df['home_scored_at_away_mean'] = df.apply(
    lambda x: x['home_scored_at_away_sum'] / x['home_scored_at_away_count'] 
    if x['home_scored_at_away_count'] > 0 else None, axis=1
)
df['away_scored_at_home_mean'] = df.apply(
    lambda x: x['away_scored_at_home_sum'] / x['away_scored_at_home_count'] 
    if x['away_scored_at_home_count'] > 0 else None, axis=1
)

# 6. 过滤掉任意一个均值无数据的行（与原代码逻辑一致）
df = df.dropna(subset=['home_scored_at_home_mean', 'away_scored_at_away_mean', 
                       'home_scored_at_away_mean', 'away_scored_at_home_mean'])

print(df)

优化说明 #

• 消除逐行循环：使用groupby+expanding实现向量式计算，时间复杂度从O(n²)降至O(n log n)，5万行数据的处理速度会大幅提升。
• 减少重复计算：批量处理分组统计，避免多次重复调用groupby，进一步优化性能。
• 逻辑清晰可读：拆分各统计项的计算逻辑，变量命名直观，便于理解和维护。
• 严谨处理边界：通过shift(1)确保只使用当前日期之前的历史数据，用dropna实现原代码中"任意场次为0则跳过"的逻辑。

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