构建回归预测函数：数据集拆分与RMSE计算问题咨询

问题1：修复isin()拆分数据集的报错 #

isin()方法要求传入列表、数组这类可迭代的类列表对象，直接传单个整数会触发报错。结合你的需求（训练集用2021年数据，测试集用其余年份），正确的拆分代码有两种写法：

写法一（用isin()）：

# 训练集：2021年数据
train = df3[df3['year'].isin([2021])]
# 测试集：非2021年的数据
test = df3[~df3['year'].isin([2021])]

写法二（直接用等于判断，更简洁）：

train = df3[df3['year'] == 2021]
test = df3[df3['year'] != 2021]

问题2：整合线性回归模型并返回RMSE #

先导入所需依赖库，再把数据集拆分、模型训练、RMSE计算整合成一个函数，以下是完整可运行代码：

import numpy as np
import math
from sklearn.linear_model import LinearRegression

def calculate_revenue_rmse(df):
    # 1. 拆分训练集和测试集
    train = df[df['year'] == 2021]
    test = df[df['year'] != 2021]
    
    # 2. 分离特征与目标变量
    # 假设特征是year和revenue之外的所有列，可根据你的实际衍生特征调整
    X_train = train.drop(['year', 'revenue'], axis=1)
    y_train = train['revenue']
    X_test = test.drop(['year', 'revenue'], axis=1)
    y_test = test['revenue']
    
    # 3. 训练最简线性回归模型
    model = LinearRegression()
    model.fit(X_train, y_train)
    
    # 4. 预测并按要求计算RMSE
    y_pred = model.predict(X_test)
    mse = np.square(np.subtract(y_test, y_pred)).mean()
    rmse = math.sqrt(mse)
    
    return rmse

# 调用函数传入你的数据框df3
final_rmse = calculate_revenue_rmse(df3)
print(f"测试集RMSE结果：{final_rmse}")

代码说明： #

• 特征列可根据实际衍生特征调整，若有不需要的特征，直接在drop()方法中添加列名即可
• LinearRegression()是sklearn库的最简线性回归模型，无需额外参数即可使用
• 严格按照你指定的MSE、RMSE计算公式实现，逻辑完全匹配需求

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