题解 | 多项式岭回归

import sys
import math
import numpy as np

def get_y(xs,ys,gidx)->float:
    lam = 0.1
    dim = 3
    y = np.array(ys)
    y = y.T
    X = np.array([
        [math.pow(e,i) for e in xs]
        for i in range(dim)
    ])
    X = X.T
    gx = np.array([math.pow(gidx,i) for i in range(dim)])
    beta = np.linalg.inv(X.T @ X + lam * np.eye(dim)) @ X.T @ y
    res = gx @ beta
    return float(res.item())



def main(M,N,gap_ids,data)->list:
    assert len(gap_ids)==M,"err1"
    L,res=0,[]
    for i,gidx in enumerate(gap_ids):
        R=N
        if i<M-1:
            R=gap_ids[i+1]
        indices = list(range(L, gidx)) + list(range(gidx + 1, R))
        xs = [i + 1 for i in indices]
        ys = [float(data[i]) for i in indices]
        y = get_y(xs,ys,gidx+1)
        L = gidx+1
        res.append(y)
    return res

if __name__ == "__main__":
    M,N = map(int, input().split())
    gap_index = []
    data_arr = []
    for i in range(N):
        input_str = input()
        if input_str.startswith('Gap'):
            gap_index.append(i)
            data_arr.append(None)
            continue
        data_arr.append(float(input_str))
    res = main(M,N,gap_index,data_arr)
    for idx,e in enumerate(res):
        print(f"Gap_{idx+1}: {e:.2f}")