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 = 1 + times[i] - 1 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)