AI003. 缺失值恢复

"""Solve the unique regression model with exact rational arithmetic."""
import sys
from fractions import Fraction


def gauss(a, rhs):
    size = len(rhs)
    for col in range(size):
        pivot = next((r for r in range(col, size) if a[r][col] != 0), None)
        if pivot is None:
            raise ValueError('The training equation must have a unique solution')
        a[col], a[pivot] = a[pivot], a[col]
        rhs[col], rhs[pivot] = rhs[pivot], rhs[col]
        div = a[col][col]
        for j in range(col, size):
            a[col][j] /= div
        rhs[col] /= div
        for r in range(size):
            if r == col:
                continue
            factor = a[r][col]
            if factor == 0:
                continue
            for j in range(col, size):
                a[r][j] -= factor * a[col][j]
            rhs[r] -= factor * rhs[col]
    return rhs


def format_four(value):
    if abs(value) < Fraction(1, 20000):
        return '0.0000'
    magnitude = abs(value)
    q, r = divmod(magnitude.numerator * 10000, magnitude.denominator)
    if 2 * r >= magnitude.denominator:
        q += 1
    sign = '-' if value < 0 and q else ''
    return f'{sign}{q // 10000}.{q % 10000:04d}'


def solve(text):
    data = iter(text.split())
    first = next(data, None)
    if first is None:
        return ''
    n, d, queries = int(first), int(next(data)), int(next(data))
    lam = Fraction(next(data))
    size = d + 1
    a = [[Fraction(0) for _ in range(size)] for _ in range(size)]
    rhs = [Fraction(0) for _ in range(size)]
    for _ in range(n):
        row = [Fraction(1)] + [Fraction(next(data)) for _ in range(d)]
        target = Fraction(next(data))
        for i in range(size):
            rhs[i] += row[i] * target
            for j in range(size):
                a[i][j] += row[i] * row[j]
    for i in range(1, size):
        a[i][i] += lam
    weights = gauss(a, rhs)
    output = []
    for _ in range(queries):
        row = [next(data) for _ in range(d)]
        target = Fraction(next(data))
        missing = row.index('?')
        coefficient = weights[missing + 1]
        if abs(coefficient) < Fraction(1, 10**10):
            output.append('NA')
            continue
        known = weights[0] + sum(weights[j+1] * Fraction(value)
                                 for j,value in enumerate(row) if j != missing)
        output.append(format_four((target - known) / coefficient))
    return '\n'.join(output)


if __name__ == '__main__':
    result = solve(sys.stdin.read())
    if result:
        print(result)