Solución del problema "A gas station too far"

A gas station too far.py

from jutge import read # Cosas del input...

def optimize(f, f_extra_args, l, r, y): 
    # Tiempo: O(t(f) log(r-l)) donde t(f) es el coste de calcular f(x)
    #
    # Este método maximiza la función f y si hay varios inputs que dan
    # el valor y encuentra el mínimo {min x | f(x) = y} y se 
    # cumple que la función f es (no estrictamente) creciente, es decir, 
    # que siempre {f(x) >= f(x+1)}.
    # Si tienes una función decreciente {f(x) <= f(x+1)} puedes convertirla 
    # en creciente cambiando el signo del resultado y así se hace un apaño
    # para no tener que modificar esta parte del código.
    # Esto a muchos les sonará como búsqueda binaria o en funciones reales
    # método de bisección :)
    min_x = None
    while l <= r:
        m = (l + r) // 2
        f_m = f(m, *f_extra_args)
        if f_m >= y:
            if f_m == y:
                min_x = m
            r = m-1 # Ya solo queremos x's menores que m
        else:
            l = m+1 # Toda x <= l cumplirá f(x) < y por lo que no nos interesan
    return min_x

def f(x, n, s, l):
    # Problema: https://jutge.org/problems/P23800_en
    # función de el número de gasolineras alcanzadas
    # respecto al tamaño de depósito (x), como es lógico, siempre
    # que el tamaño del depósito sea mayor, llegaremos a la misma gasolinera
    # o a una más lejana, por lo que deducimos que esta función es monótona, por lo que
    # podemos optimizarla con optimize()
    refill_cnt = 0
    fuel = x
    for i in range(n):
        if x < l[i]: # imposible pasar de esta con este tamaño de depósito
            return i
        if fuel < l[i]:
            if refill_cnt == s: # todas las recargas han sido usadas, no podemos avanzar
                return i
            refill_cnt += 1
            fuel = x # regargar depósito
        fuel -= l[i] # consumir combustible correspondiente
    return n # Hemos pasado todas las gasolineras, resultado óptimo


def main():
    n = read(int)
    while n != None: # Cosas del input...
        s = read(int)
        l = [read(int) for i in range(n)]
        print(optimize(f, [n, s, l], 0, 10**9, n))
        n = read(int)


main()