Advent of Code, día 10

Recursividad y memoization, y un buen rato dándome golpes contra la pared

Voy directamente al grano:

Parte 1

Facilísima… huele a chamusquina

def solve(input):
    adapters = [0] + sorted(
                         int(n) for n in input.strip().split('\n'))
    counts = {
        1: 0,
        2: 0,
        3: 0,
    }
    for i in range(1, len(adapters)):
        diff = adapters[i] - adapters[i - 1]
        counts[diff] += 1

    return counts[1] * (1 + counts[3])

Parte 2

En mi cabeza, la idea de lo que quería hacer estaba clarísima, pero mi código se negaba a calcular las respuestas correctas para los casos de prueba.

La idea: algoritmo recursivo. Voy recorriendo el array de adaptadores ordenados. En un cada punto, el número de ordenaciones válidas es la suma de las ordenaciones válidas del resto de los adaptadores a) suponiendo que no me salto ningún adaptador, b) suponiendo que me salto uno o c) suponiendo que me salto dos. Si al saltarme uno o dos adaptadores, la diferencia es demasiado grande, el número de ordenaciones válidas de esa variante es 0.

También tenía en la cabeza que ese algoritmo de “fuerza bruta” iba a ser muy ineficiente, ya que examina repetidamente muchas ramas del árbol de posibilidades, así que iba a tener que usar “memoization” (luego lo cuento)

Tras un buen rato de depurar, inventarme casos de prueba sencillos y resolverlos con papel y boli, poner prints… conseguí ajustar bien todos los casos recursivos y obtener la respuesta correcta. El código final parece sencillo, pero no hace justicia al tiempo y esfuerzo que me ha llevado escribirlo… ¡menudo calvario!

La solución:

def valid_arrangements(adapters, start):
    if start + 1 == len(adapters):
        return 1

    arrangements = 0
    for i in [1, 2, 3]:
        if start + i < len(adapters) and \
           adapters[start + i] - adapters[start] <= 3:
            arrangements += valid_arrangements(adapters, start + i)

    return arrangements
def solve(input):
    adapters = sorted(int(n) for n in input.strip().split('\n'))
    adapters = [0] + adapters + [adapters[-1] + 3]
    return valid_arrangements(adapters, 0)

Efectivamente, sin memoización la cosa no pinta bien… la complejidad es O(2ⁿ) y si ponemos este código a funcionar con n=100 nos podemos aburrir de esperar… he hecho unas gráficas con n entre 20 y 45. Con n=45 el tiempo de ejecución ya es de mas de 1,5 minutos

Tiempos de ejecución SIN memoización

¿Qué es eso de la memoización?

Si nos fijamos en este código, exploramos las mismas ramas del árbol de combinaciones una y otra vez. La memoización no es más que cachear los resultados ya calculados para evitar la mayoría de llamadas recursivas. Tiene la ventaja de que acelera la ejecución del algoritmo de O(2ⁿ) a O(n) sin necesidad de cambiar el código principal.

En Python podemos “cachear” los resultados de funciones aplicando un decorador a la función:

@memoize
def valid_arrangements(adapters, start):
    # ... el resto del código no cambia ...

El decorador sirve para envolver una función dentro de otra. Esta es la definición del decorador:

def memoize(f):
    memo = {}

    @wraps(f)
    def helper(adapters, start):
        key = (tuple(adapters[:10]), start)
        if key not in memo:
            memo[key] = f(adapters, start)
        return memo[key]

    return helper

La función memoize se aplica a la función valud_arrangements, y la “envuelve” con la función helper(). Una vez aplicado el decorador, cuando llamemos a valid_arrangements() en realidad estamos llamando a la función helper() que nos ha devuelto el decorador.

La función helper() comprueba si ya nos han llamado con los mismos argumentos antes. Si es que si, devuelve el resultado que tiene en su memoria interna (el dict memo). Si es que no, llama a la función real y a la vuelta, almacena el resultado en memo.

Como uno de nuestros parámetros es un array de 100 números, si usáramos los parámetros completos la clave de la caché sería muy larga, y por cada llamada estaríamos comparando los 100 números de la entrada una y otra vez. Para aligerar esto un poco, construyo la clave usando solo las primeras posiciones (porque sabemos que el array no cambia).

El @wraps(f) es un decorador built-in que se aplica a la función interna para que la función resultante envuelta tenga las mismas propiedades (p.ej. la propiedad __doc__) que la función original (más info aquí)

En Python 3.9 ya tenemos un decorador built-in llamado cache que podría haber usado en lugar de escribirme mi propio memoize… pero lo he visto tarde! 🤦 También hay otro decorador llamado lru_cache que es igual pero con un límite de capacidad (ese lo había usado alguna vez pero no lo recordaba).

¿Funciona?

Entonces… esto del memoize… ¿funciona? ¡Y tanto que si! Se obtiene la respuesta correcta en menos de 1ms. Estrella dorada al saco y otro día completado.

Otras soluciones

No he tenido tiempo de ver otras soluciones en detalle, pero por lo visto, la entrada tenía truco (solo había diferencias de 1 o de 3 entre los adaptadores, y no había secuencias muy largas de 1 seguidos) por lo que el problema se podía simplificar muchísimo.

Rubén Valseca se ha basado en esa simplificación en su solución

@mariomka lo ha resuelto en Rust, “acumulando de cuantas formas puedes llegar a un nodo”, y devolviendo el número de formas de las que puedes llegar al último nodo. El código es bastante sencillo y eficiente O(n) (si no contamos el O(n log n) de la ordenación inicial necesaria)

Mi solución completa aquí: https://github.com/jabadia/adventOfCode2020/tree/main/d10