Tidy up day 13 solution

day-13/day-13.py

@@ -2,17 +2,17 @@
 from pathlib import Path
 
 
-def fold(i, dots):
+def fold(axis, fold_at, dots):
     new_dots = set()
-    for d in dots:
-        if i[0] == 'x':
-            n = (i[1] - abs(d[0] - i[1]), d[1])
-            if d[0] != i[1]:
-                new_dots.add(n)
-        elif i[0] == 'y':
-            n = (d[0], i[1] - abs(d[1] - i[1]))
-            if d[1] != i[1]:
-                new_dots.add(n)
+    for (x, y) in dots:
+        if axis == 'x':
+            new_x = fold_at - abs(x - fold_at)
+            if new_x != fold_at:
+                new_dots.add((new_x, y))
+        elif axis == 'y':
+            new_y = fold_at - abs(y - fold_at)
+            if new_y != fold_at:
+                new_dots.add((x, new_y))
     return new_dots
 
 
@@ -27,7 +27,8 @@ def part_1(input):
         elif 'fold along' in line:
             d, c = line.strip().split()[-1].split('=')
             instructions.append((d, int(c)))
-    dots = fold(instructions.pop(0), dots)
+    axis, fold_at = instructions.pop(0)
+    dots = fold(axis, fold_at, dots)
     result = len(dots)
     print("Part 1 result:", result)
 
@@ -43,14 +44,11 @@ def part_2(input):
         elif 'fold along' in line:
             d, c = line.strip().split()[-1].split('=')
             instructions.append((d, int(c)))
-
     while len(instructions):
-        dots = fold(instructions.pop(0), dots)
-    max_x = 0
-    max_y = 0
-    for d in dots:
-        max_x = max(max_x, d[0] + 1)
-        max_y = max(max_y, d[1] + 1)
+        axis, fold_at = instructions.pop(0)
+        dots = fold(axis, fold_at, dots)
+    max_x = max([x for x, _ in dots])
+    max_y = max([y for _, y in dots])
     print("Part 2 result:")
     for y in range(max_y):
         for x in range(max_x):