AoC/aoc2023/day6/day6.py

import os
from math import sqrt, ceil, floor

with open(os.path.join(os.path.dirname(__file__), "example.txt")) as example:
    example_data = example.read().splitlines()

with open(os.path.join(os.path.dirname(__file__), "input.txt")) as example:
    input_data = example.read().splitlines()


def solve(times, dists):
    def root_eq(b, c):
        D = b**2 - 4 * c
        x1 = (b - sqrt(D)) / 2
        x2 = (b + sqrt(D)) / 2
        return x1, x2

    total = 1
    for time, dist in zip(times, dists):
        at_least, less_then = root_eq(time, dist)
        solution_number = ceil(less_then) - floor(at_least) - 1
        total *= solution_number
    print(total)
    return total


def parse1(data):
    times = [int(n) for n in data[0][5:].split()]
    dists = [int(n) for n in data[1][9:].split()]
    return times, dists


def parse2(data):
    times = [int(data[0][5:].replace(" ", ""))]
    dists = [int(data[1][9:].replace(" ", ""))]
    return times, dists


solve(*parse1(example_data))
solve(*parse1(input_data))
solve(*parse2(example_data))
solve(*parse2(input_data))