from math import sqrt, ceil, floor from utils import parse_day from collections import defaultdict tx_min, tx_max, ty_min, ty_max = \ map(int, parse_day(17, r'.* x=(-?\d+)..(-?\d+), y=(-?\d+)..(-?\d+)')) vx_min = int(sqrt(tx_min * 2) - 1) vx_max = tx_max vy_min = ty_min vy_max = abs(ty_min) - 1 def sum_n(n): return n * (n + 1) // 2 print(sum_n(vy_max)) def time_at_pos(v, d): """d = -t(t - 2v - 1) / 2 solved for t""" middle = sqrt((2 * v + 1) ** 2 - 8 * d) return (-middle + 2 * v + 1) / 2, (middle + 2 * v + 1) / 2 def time_at_xpos(v, d): if d > sum_n(v): return float('inf') return time_at_pos(v, d)[0] def y_times(vy0): start = ceil(time_at_pos(vy0, ty_max)[1]) end = floor(time_at_pos(vy0, ty_min)[1]) return start, end def x_times(vx0): start = time_at_xpos(vx0, tx_min) end = time_at_xpos(vx0, tx_max) return start, end times = defaultdict(set) for vy0 in range(vy_min, vy_max + 1): tmin, tmax = y_times(vy0) for t in range(tmin, tmax + 1): times[t].add(vy0) points = set() for vx0 in range(vx_min, vx_max + 1): tmin, tmax = x_times(vx0) for t, vy0set in times.items(): if tmin <= t <= tmax: for vy0 in vy0set: points.add((vx0, vy0)) print(len(points))