day 17: faster again

1 files changed, 6 insertions, 6 deletions

diff --git a/17.py b/17.py
index be5bccb..ebe3116 100644
--- a/17.py
+++ b/17.py
@@ -1,6 +1,7 @@
 from math import sqrt, ceil, floor
 from utils import parse_day
 from collections import defaultdict
+from itertools import chain
 
 tx_min, tx_max, ty_min, ty_max = \
     map(int, parse_day(17, r'.* x=(-?\d+)..(-?\d+), y=(-?\d+)..(-?\d+)'))
@@ -39,11 +40,10 @@ 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()
+count = 0
 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))
+    count += len(set(chain.from_iterable(
+        vy0set for t, vy0set in times.items() if tmin <= t <= tmax
+    )))
+print(count)