advent-of-code/2020/solutions/13/python/main.py

from os.path import join, dirname
from typing import List, Tuple
from math import ceil, floor


class Station:

    def __init__(self, txt: str) -> None:
        arrival, busses = txt.strip().split('\n')
        self.arrival = int(arrival)
        self.busses = [
            int(bus)
            for bus in busses.replace('x', '0').split(',')
        ]

    @staticmethod
    def get_next_for_id(id: int, arrival: int) -> int:
        return id * ceil(arrival / id)

    def find_next(self) -> Tuple[int, int]:
        arrivals: List[Tuple[int, int]] = []
        for bus in self.busses:
            if bus == 0:
                continue
            arrivals.append((bus, self.get_next_for_id(bus, self.arrival)))
        return min(arrivals, key=lambda x: x[1])

    def get_flag(self) -> int:
        id, arrives = self.find_next()
        return id * (arrives - self.arrival)

    def contest(self, offset: int = 1) -> int:
        # Prepare
        highest = max(self.busses)
        highest_i = self.busses.index(highest)
        others = [
            (bus, i - highest_i)
            for i, bus in enumerate(self.busses)
            if bus != 0 and i != highest_i
        ]
        others.sort(key=lambda x: x[0], reverse=True)

        # Compute
        i: int = max(1, floor(offset / highest))
        while True:
            x = highest * i
            error = False
            for bus, diff in others:
                dx = x + diff
                if dx != self.get_next_for_id(bus, dx):
                    error = True
                    break
            if not error:
                return x - highest_i
            i += 1


data = join(dirname(__file__), '../data.txt')
with open(data) as f:

    # Some "testing"
    all = {
        '17,x,13,19': 3417,
        '67,7,59,61': 754018,
        '67,x,7,59,61': 779210,
        '67,7,x,59,61': 1261476,
        '1789,37,47,1889': 1202161486,
    }
    for busses, expected in all.items():
        station = Station('69\n' + busses)
        print(expected, expected == station.contest())

    txt = f.read()
    station = Station(txt)
    print(station.get_flag())
    print(station.contest(offset=10**14))