如何从自定义线性回归函数输出为Pandas DataFrame添加新列

我来帮你搞定这个线性回归预测值添加的问题！你的核心需求是把自己计算的回归预测值y_hat加到原始DataFrame里，或者生成独立的预测值列表，咱们一步步来解决：

首先修正你的函数错误 #

你原来的函数里没有定义mean_X和mean_Y，这会直接导致运行报错。另外，直接把df写在函数里不够灵活，最好让函数接收原始DataFrame和列名，这样更通用。

方案1：让函数直接返回带预测列的DataFrame #

这个方案最省心，函数会自动从DataFrame提取数据、计算回归参数、生成预测值并添加新列：

def add_prediction(df, x_column, y_column):
    # 从DataFrame提取X和Y的列表
    X = df[x_column].tolist()
    Y = df[y_column].tolist()
    
    # 计算样本均值
    mean_X = sum(X) / len(X)
    mean_Y = sum(Y) / len(Y)
    
    # 计算回归系数beta1和beta0
    numerator = sum([(X[i] - mean_X) * (Y[i] - mean_Y) for i in range(len(X))])
    denominator = sum([(X[i] - mean_X)**2 for i in range(len(X))])
    beta1 = numerator / denominator
    beta0 = mean_Y - beta1 * mean_X
    
    # 生成预测值列表
    y_hat = [beta0 + beta1 * x for x in X]
    
    # 添加新列并返回新的DataFrame（原DataFrame不会被修改）
    return df.assign(prediction=y_hat)

# 使用示例：假设你的df有'X'和'Y'两列
df_with_pred = add_prediction(df, 'X', 'Y')

方案2：单独生成预测值列表再手动添加 #

如果你想在函数外部控制预测值的使用，可以单独写一个计算y_hat的函数，之后手动加到DataFrame里：

def calculate_prediction(X, Y):
    mean_X = sum(X) / len(X)
    mean_Y = sum(Y) / len(Y)
    
    numerator = sum([(X[i] - mean_X) * (Y[i] - mean_Y) for i in range(len(X))])
    denominator = sum([(X[i] - mean_X)**2 for i in range(len(X))])
    beta1 = numerator / denominator
    beta0 = mean_Y - beta1 * mean_X
    
    return [beta0 + beta1 * x for x in X]

# 使用示例：从df提取X和Y列表
X = df['X'].tolist()
Y = df['Y'].tolist()
y_hat = calculate_prediction(X, Y)

# 手动添加到原DataFrame
df['prediction'] = y_hat

为什么numpy.insert没成功？ #

numpy.insert的作用是往numpy数组里插入元素，但你的需求是给DataFrame添加列，两者的逻辑不一样：

• 如果你把X转成数组后用numpy.insert插入y_hat，得到的是一个新的二维数组，和原DataFrame的索引、列名等结构完全脱节，还要重新构造DataFrame，非常麻烦。
• 直接给DataFrame赋值新列（df['prediction'] = y_hat）是最直接的方式，只要y_hat的长度和df的行数一致，就能自动匹配。

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