如何优化Pandas按多列分组执行岭回归的嵌套循环实现

你可以通过Pandas内置的groupby+apply逻辑完全替换三层嵌套循环，核心优化思路是把原来每个分组内单独执行的聚合操作提前到全局做，再批量处理所有分组的回归任务，性能会有数量级的提升。

完整实现代码如下：

import pandas as pd
import numpy as np
from sklearn.linear_model import Ridge

# 第一步：全局预聚合，一次性完成所有[A,B,C,q]组合的value求和，避免循环内重复计算
agg_df = df.groupby(['A', 'B', 'C', 'q'], as_index=False)['value'].sum()

# 第二步：定义单分组的岭回归处理函数，可自定义返回需要的回归结果
def ridge_regression_per_group(group_df):
    # 样本量不足2条时无法拟合回归，可根据需求自定义兜底逻辑
    if len(group_df) < 2:
        group_df['value_pred'] = group_df['value']
        group_df['ridge_coef'] = np.nan
        group_df['ridge_intercept'] = np.nan
        return group_df
    # 拟合岭回归，alpha参数可根据实际需求调整
    X = group_df['q'].values.reshape(-1, 1)
    y = group_df['value'].values
    ridge = Ridge(alpha=1.0)
    ridge.fit(X, y)
    # 写入预测值、回归系数、截距等结果，可按需增减
    group_df['value_pred'] = ridge.predict(X)
    group_df['ridge_coef'] = ridge.coef_[0]
    group_df['ridge_intercept'] = ridge.intercept_
    return group_df

# 第三步：按A/B/C分组批量应用回归函数
result_df = agg_df.groupby(['A', 'B', 'C'], group_keys=False).apply(ridge_regression_per_group)

该方案的性能优势来源于两个核心优化：

• 全局聚合是Pandas底层优化的向量化操作，远快于循环内每个分组单独做查询聚合
• groupby.apply的分组遍历逻辑经过Pandas内部优化，避免了三层循环中重复的query查询开销

如果只需要提取每个分组的回归参数不需要预测值，可修改处理函数返回字典，最后用result_df = agg_df.groupby(['A', 'B', 'C']).apply(ridge_regression_per_group).reset_index()得到每个分组一行的回归参数表。

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