题解 | 小红的语言模型推理耗时预测 o(m+n)

简单说个思路吧，每轮都使用全量数据迭代，数据的三个特征从头到尾都不变，可以尝试把计算公式写出来展开，参数提到Σ之前，发现要循环m次计算的三个项其实是恒定不变的，每次迭代无需循环m次计算，具体可以自己写出来思考下，还是挺有意思的。

import sys

input_data = sys.stdin.read().split()
if not input_data:
    sys.exit()

m = int(input_data[0])
n = int(input_data[1])
a = float(input_data[2])

f1 = [0.0] * m
f2 = [0.0] * m
f3 = [0.0] * m
y = [0.0] * m

idx = 3
for i in range(m):
    f1[i] = int(input_data[idx])
    f2[i] = int(input_data[idx+1])
    f3[i] = int(input_data[idx+2])
    y[i] = int(input_data[idx+3])
    idx += 4

max1, min1 = max(f1), min(f1)
max2, min2 = max(f2), min(f2)
max3, min3 = max(f3), min(f3)

d1 = max1 - min1
if d1 == 0:
    f1 = [0.0] * m
    d1 = 1.0
else:
    f1 = [(x - min1) / d1 for x in f1]

d2 = max2 - min2
if d2 == 0:
    f2 = [0.0] * m
    d2 = 1.0
else:
    f2 = [(x - min2) / d2 for x in f2]

d3 = max3 - min3
if d3 == 0:
    f3 = [0.0] * m
    d3 = 1.0
else:
    f3 = [(x - min3) / d3 for x in f3]

S1 = sum(f1)
S2 = sum(f2)
S3 = sum(f3)
S11 = sum(x * x for x in f1)
S22 = sum(x * x for x in f2)
S33 = sum(x * x for x in f3)
S12 = sum(f1[i] * f2[i] for i in range(m))
S13 = sum(f1[i] * f3[i] for i in range(m))
S23 = sum(f2[i] * f3[i] for i in range(m))
Sy = sum(y)
Sy1 = sum(y[i] * f1[i] for i in range(m))
Sy2 = sum(y[i] * f2[i] for i in range(m))
Sy3 = sum(y[i] * f3[i] for i in range(m))

w1 = w2 = w3 = b = 0.0

for _ in range(n):
    g0 = (w1 * S1 + w2 * S2 + w3 * S3 + b * m - Sy) / m
    g1 = (w1 * S11 + w2 * S12 + w3 * S13 + b * S1 - Sy1) / m
    g2 = (w1 * S12 + w2 * S22 + w3 * S23 + b * S2 - Sy2) / m
    g3 = (w1 * S13 + w2 * S23 + w3 * S33 + b * S3 - Sy3) / m
    
    b -= a * g0
    w1 -= a * g1
    w2 -= a * g2
    w3 -= a * g3

if max1 != min1:
    w1 /= d1
if max2 != min2:
    w2 /= d2
if max3 != min3:
    w3 /= d3

b = b - w1 * min1 - w2 * min2 - w3 * min3

print(f"{b:.2f} {w1:.2f} {w2:.2f} {w3:.2f}")