summaryrefslogtreecommitdiffstats
path: root/17.py
blob: be5bccb23a36776f5ccdc9083dbd1be882504b90 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
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))