牛客题解 | 计算回归分析的R-squared

题目

题目链接

R-squared是一种衡量回归模型性能的指标，其计算公式为：
$$R^2=1-\frac{SSR}{SST}$$
其中，SSR是残差平方和，SST是总平方和。

标准代码如下

def r_squared(y_true, y_pred):
    if np.array_equal(y_true, y_pred):
        return 1.0

    y_mean = np.mean(y_true)

    ssr = np.sum((y_true - y_pred) ** 2)

    sst = np.sum((y_true - y_mean) ** 2)

    try:
        r2 = 1 - (ssr / sst)
        if np.isinf(r2):
            return 0.0
        return round(r2, 3)
    except ZeroDivisionError:
        return 0.0

在sklearn库中也有对应的计算方法,可以使用metrics模块中的r2_score函数:

def r_squared(y_true, y_pred):
    from sklearn.metrics import r2_score
    r2 = r2_score(y_true, y_pred)
    return round(float(r2), 3)