AoC2023/day06/task1.py

from math import sqrt, floor, ceil


def do_task(input):
    lines = input.split("\n")

    times = [int(t) for t in lines[0].split()[1:]]
    distances = [int(d) for d in lines[1].split()[1:]]

    prod = 1

    for i in range(len(times)):
        a = -1
        b = times[i]
        c = -distances[i]

        max_h = (-b - sqrt(b**2 - 4 * a * c)) / (2 * a)
        min_h = (-b + sqrt(b**2 - 4 * a * c)) / (2 * a) + 1

        prod *= ceil(max_h) - floor(min_h)

    return prod

# This actually works
def do_task_the_dumb_way(input):
    lines = input.split("\n")

    times = [int(t) for t in lines[0].split()[1:]]
    distances = [int(d) for d in lines[1].split()[1:]]

    prod = 1

    for i in range(len(times)):
        ways = 0
        for h in range(times[i]):
            d = 0
            for t in range(times[i] - h):
                d += h
            if d > distances[i]:
                ways += 1
        prod *= ways

    return prod


with open("day06/data/input100.txt") as infile:
    result = do_task(infile.read())
    print(result)
    print(result == 288)

with open("day06/data/input101.txt") as infile:
    result = do_task(infile.read())
    print(result)
    print(result == 227850)