aoc day 7 linear time solution

import sys
import numpy as np

def compute_fuel(x, a, b, c):
    return a*x**2 + b*x + c

def main(input_file=sys.stdin, output_file=sys.stdout):
    # Read input vector.
    ps = np.array([int(e) for e in input_file.readline().split(',')], dtype=np.int)

    ps.sort()

    # Define the fuel cost function fuel(x) as cost for all crabs to get to x
    # fuel(x) is piece-wise quadratic between all the crab_positions p
    # fuel(x) = sum (abs(p-x) * (abs(p-x) + 1)) / 2 over all p, with p the crab_positions

    # In left-most piece of fuel(x) all crab positions p are right of x, so abs(p-x) = p-x

    # fuel(x) = ax^2 + bx + c
    a = len(ps) / 2
    b = (-2*ps.sum() - len(ps)) / 2
    c = ((ps*ps).sum()+ps.sum()) / 2

    # Definitely not optimal in first and last quadratic piece

    # Compute fuel at first p
    best_fuel = compute_fuel(ps[0], a, b, c)
    best_fuel_point = ps[0]

    # Every time we cross a crab position fuel(x) is increased by x-p
    for p, p_next in zip(ps[:-1], ps[1:]):
        b += 1
        c -= p

        # Either optimum is between p and p_next:
        # In which case compute min of quadratic piece

        # Derivative of current quadratic piece equal to zero at:
        x_opt = -b / (2*a)
        if p <= x_opt <= p_next:
            # Point needs to be integer
            x_opt_ceil = np.ceil(x_opt)
            x_opt_floor = np.floor(x_opt)

            fuel_at_x_opt_ceil = compute_fuel(x_opt_ceil, a, b, c)
            if fuel_at_x_opt_ceil < best_fuel:
                best_fuel = fuel_at_x_opt_ceil
                best_fuel_point = x_opt_ceil

            fuel_at_x_opt_floor = compute_fuel(x_opt_floor, a, b, c)
            if fuel_at_x_opt_floor < best_fuel:
                best_fuel = fuel_at_x_opt_floor
                best_fuel_point = x_opt_floor

        # Or optimum is at p_next
        fuel_at_p_next = compute_fuel(p_next, a, b, c)
        if fuel_at_p_next < best_fuel:
            best_fuel = fuel_at_p_next
            best_fuel_point = p_next

    print(f"{best_fuel} at {best_fuel_point}")

    output_file.write(str(int(best_fuel)))

if __name__ == "__main__":
    main()