如何高效对DataFrame分组子集应用线性回归并生成预测结果

需要按V1、V2、V3、V4、V5、V6对原始数据分组，每组应用线性回归预测值，最终生成包含V1、V2、V3、V4、V5、V6、time、Predicted value的DataFrame。当前实现返回的对象难以进一步处理，求高效实现方案。

现有代码 #

def model(df):     
    X = df['time'].to_numpy().reshape((-1, 1))
    Y = df['speed'].to_numpy() 
    X_new = np.arange(1, 60, 1).reshape((-1, 1))
    return np.squeeze(LinearRegression().fit(X, Y).predict(X_new))

def group_predictions(df): 
    return df.groupby(['V1', 'V2', 'V3', 'V4', 'V5','V6']).apply(model)

解决方案 #

核心调整方向 #

修改model函数，让它返回包含time和预测值的结构化DataFrame，而非单纯的数值数组；再通过reset_index()将分组键从索引转为普通列，直接得到目标格式的结果。

修改后的完整代码 #

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

def model(df):     
    X = df['time'].to_numpy().reshape((-1, 1))
    Y = df['speed'].to_numpy() 
    # 生成要预测的time序列
    forecast_time = np.arange(1, 60, 1)
    # 拟合模型并计算预测值
    pred_values = LinearRegression().fit(X, Y).predict(forecast_time.reshape(-1, 1))
    # 返回带time和预测值的DataFrame
    return pd.DataFrame({
        'time': forecast_time,
        'Predicted value': pred_values
    })

def group_predictions(df): 
    # 分组应用模型后，重置索引将分组键转为列
    result = df.groupby(['V1', 'V2', 'V3', 'V4', 'V5','V6']).apply(model)
    return result.reset_index(drop=False).drop(columns='level_6')

关键说明 #

• model函数返回的DataFrame直接包含预测对应的time和Predicted value，保证每组结果都是结构化的表格，避免了后续数组转结构的麻烦。
• groupby.apply执行后，分组键（V1到V6）会成为结果的多级索引，调用reset_index()可以把这些索引转为普通列；最后删除自动生成的level_6列（分组后每组内部的索引），就得到完全符合需求的DataFrame。

可选优化（针对大数据量） #

如果处理的数据集非常大，可以考虑用dask进行并行分组处理，或者使用statsmodels的公式接口简化回归代码，但上述方案已经足够简洁高效，覆盖绝大多数常规场景。

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