advent-of-code/2023/day6.py

import math
import sys


def extreme_times_to_beat(time: int, distance: int) -> tuple[int, int]:
    # formula (T = total time, D = target distance)
    # - speed(t) = t,
    # - distance(t) = (T - t) * speed(t)
    # - distance(t) > D
    #   <=> (T - t) * t > D
    #   <=> -t^2 + T * t - D >= 0

    a, b, c = -1, time, -distance
    d = b * b - 4 * a * c

    t1 = math.ceil(-(b - math.sqrt(d)) / (2 * a))
    t2 = math.floor(-(b + math.sqrt(d)) / (2 * a))

    if (time - t1) * t1 == distance:
        t1 += 1

    if (time - t2) * t2 == distance:
        t2 -= 1

    return t1, t2


lines = sys.stdin.read().splitlines()

# part 1
times = list(map(int, lines[0].split()[1:]))
distances = list(map(int, lines[1].split()[1:]))
answer_1 = math.prod(
    t2 - t1 + 1
    for t1, t2 in (
        extreme_times_to_beat(time, distance)
        for time, distance in zip(times, distances)
    )
)
print(f"answer 1 is {answer_1}")

# part 2
time = int(lines[0].split(":")[1].strip().replace(" ", ""))
distance = int(lines[1].split(":")[1].strip().replace(" ", ""))
t1, t2 = extreme_times_to_beat(time, distance)
answer_2 = t2 - t1 + 1
print(f"answer 2 is {answer_2}")