Day 05 (glamorized)

day05/part2.py

@@ -1,26 +1,18 @@
 #!/usr/bin/env python3
 import re
 from collections import defaultdict
+from itertools import repeat
 
 lines = (x.strip() for x in open("input.txt"))
 points = defaultdict(int)
 
 for line in lines:
     [x1, y1, x2, y2] = list(map(int, re.match(r"(\d+),(\d+) -> (\d+),(\d+)", line).groups()))
-    if x1 == x2:
-        for y in range(min(y1, y2), max(y1, y2) + 1):
-            key = '{},{}'.format(x1, y)
-            points[key] = points[key] + 1
-    elif y1 == y2:
-        for x in range(min(x1, x2), max(x1, x2) + 1):
-            key = '{},{}'.format(x, y1)
-            points[key] = points[key] + 1
-    else:
-        for i in range(0, abs(x1 - x2) + 1):
-            key = '{},{}'.format(
-                (x1 + i) if x1 <= x2 else (x1 - i),
-                (y1 + i) if y1 <= y2 else (y1 - i)
-            )
-            points[key] = points[key] + 1
+    for x, y in zip(
+            range(min(x1, x2), max(x1, x2) + 1) if x1 != x2 else repeat(x1),
+            range(min(y1, y2), max(y1, y2) + 1) if y1 != y2 else repeat(y1)
+    ):
+        key = '{},{}'.format(x, y)
+        points[key] = points[key] + 1
 
-print(sum(1 if val >= 2 else 0 for val in points.values()))
+print(sum(val >= 2 for val in points.values()))