如何高效为DataFrame生成新列：基于行过滤且规避iterrows()

问题背景 #

我有一个为现有DataFrame生成新列的函数，当前用iterrows()实现，曾尝试按Year和Week分组用.apply()优化，但提速不明显。想知道怎么用向量化函数更高效完成任务。

示例数据 #

输入数据（input_df） #

历史数据（historical_df） #

计算逻辑 #

对input_df中的每一行：

• 基于Year、Week和Company过滤historical_df，仅保留当前公司时间早于当前行的历史数据
• 将过滤后的历史数据按['Year', 'Week']降序排序，取最近16周的数据
• 用这最多16周的Volume计算变异系数（CV=标准差/均值），数据量不足2条或均值为0时返回None
• 将结果存入input_df的新列Variability

现有低效实现 #

def add_variability(input_df, historical_df):
    variability_calcs = []

    for index, row in input_df.iterrows():
        filtered_historical_df = historical_df[
            ((historical_df['Year'] < row['Year']) |
            ((historical_df['Year'] == row['Year']) & (historical_df['Week'] < row['Week']))) & 
            (historical_df['Company'] == row['Company'])
        ]
    
        filtered_historical_df = filtered_historical_df.sort_values(by=['Year', 'Week'], ascending=False)
        recent_weeks = filtered_historical_df.head(16)

        if recent_weeks.shape[0] > 1:
            if recent_weeks['Volume'].mean() != 0:
                cv = recent_weeks['Volume'].std() / recent_weeks['Volume'].mean()
            else:
                cv = None
        else:
            cv = None
            
        variability_calcs.append(cv)
    
    input_df['Variability'] = variability_calcs
    return input_df

向量化优化方案 #

核心思路 #

1. 统一时间索引：将Year+Week转换为连续的total_weeks数值（如Year*52 + Week），简化时间先后的数值比较
2. 分组批量处理：按Company分组，避免逐行过滤整个历史数据集
3. 累加和快速计算窗口统计：用numpy的累加和（cumsum）一次性计算所有窗口的均值、方差，替代多次切片子DataFrame的低效操作

实现代码 #

import pandas as pd
import numpy as np

def add_total_weeks(df):
    """将Year和Week转换为连续的总周数，方便时间比较"""
    df['total_weeks'] = df['Year'] * 52 + df['Week']
    return df

def calculate_cv(group):
    """按Company分组计算每个输入行的变异系数"""
    # 拆分组内的输入行和历史数据，按时间排序
    input_rows = group[group['is_input']].sort_values('total_weeks').reset_index(drop=True)
    hist_rows = group[~group['is_input']].sort_values('total_weeks').reset_index(drop=True)
    
    if hist_rows.empty:
        input_rows['Variability'] = None
        return input_rows
    
    # 用searchsorted快速找到每个输入行对应的历史数据截止位置
    hist_weeks = hist_rows['total_weeks'].values
    input_weeks = input_rows['total_weeks'].values
    indices = np.searchsorted(hist_weeks, input_weeks)
    
    # 计算每个输入行对应的16周窗口的起始、结束索引
    starts = np.maximum(0, indices - 16)
    ends = indices
    
    # 预计算历史Volume的累加和、累加平方和，用于快速计算窗口统计量
    hist_volume = hist_rows['Volume'].values
    cum_sum = np.cumsum(hist_volume)
    cum_sum_sq = np.cumsum(hist_volume ** 2)
    
    cv_list = []
    for s, e in zip(starts, ends):
        window_size = e - s
        if window_size < 2:
            cv_list.append(None)
            continue
        
        # 计算窗口内的总和、平方和
        if s == 0:
            window_sum = cum_sum[e-1]
            window_sum_sq = cum_sum_sq[e-1]
        else:
            window_sum = cum_sum[e-1] - cum_sum[s-1]
            window_sum_sq = cum_sum_sq[e-1] - cum_sum_sq[s-1]
        
        mean = window_sum / window_size
        if mean == 0:
            cv_list.append(None)
            continue
        
        # 计算方差和标准差
        var = (window_sum_sq / window_size) - (mean ** 2)
        std = np.sqrt(var)
        cv_list.append(std / mean)
    
    input_rows['Variability'] = cv_list
    return input_rows

# 预处理：添加总周数标记
input_df = add_total_weeks(input_df.copy())
historical_df = add_total_weeks(historical_df.copy())

# 标记输入/历史数据，合并数据集
input_df['is_input'] = True
historical_df['is_input'] = False
combined_df = pd.concat([input_df, historical_df], ignore_index=True)

# 按Company分组计算，整理最终结果
result = combined_df.groupby('Company').apply(calculate_cv).reset_index(drop=True)
result = result.drop(columns=['total_weeks', 'is_input'])

print(result)

优化效果说明 #

• 时间复杂度从原方法的O(N*M)（N为输入行数，M为历史数据行数）降至O((N+M)log(N+M))（主要来自排序操作），数据量越大提速越明显
• 避免了iterrows()的逐行循环和重复过滤，利用numpy底层的向量化操作（searchsorted、cumsum）大幅提升效率
• 每个Company仅排序一次历史数据，减少重复计算

预期输出 #

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