Implement 2024 day 20

2024/src/aoc/days/day20.py

@@ -0,0 +1,82 @@
+import itertools
+
+import numpy
+
+from . import SeparateRunner
+
+DIRECTIONS = [
+    (-1, 0),
+    (1, 0),
+    (0, -1),
+    (0, 1),
+]
+
+CHEATS = [
+    (-2, 0),
+    (2, 0),
+    (0, -2),
+    (0, 2),
+]
+
+
+def parse_path(input: str) -> dict[tuple[int, int], int]:
+    grid = numpy.array(list(map(list, input.strip().split("\n"))))
+
+    ys, xs = numpy.nonzero(grid == "S")
+    sx, sy = int(xs[0]), int(ys[0])
+
+    nx, ny = sx, sy
+
+    path = {
+        (sx, sy): 0,
+    }
+
+    while grid[ny, nx] != "E":
+        x, y = nx, ny
+
+        for dx, dy in DIRECTIONS:
+            if grid[y + dy, x + dx] == "#" or (x + dx, y + dy) in path:
+                continue
+            nx = x + dx
+            ny = y + dy
+            break
+
+        path[nx, ny] = len(path)
+    return path
+
+
+def get_savings(a: tuple[tuple[int, int], int], b: tuple[tuple[int, int], int]) -> int:
+    (ax, ay), ad = a
+    (bx, by), bd = b
+
+    dist = abs(bx - ax) + abs(by - ay)
+    if dist <= 20:
+        return bd - ad - dist
+    else:
+        return 0
+
+
+class DayRunner(SeparateRunner):
+    @classmethod
+    def part1(cls, input: str, limit: int = 100) -> int:
+        path = parse_path(input)
+
+        total = 0
+
+        for (px, py), dist in path.items():
+            for dx, dy in CHEATS:
+                if (other := path.get((px + dx, py + dy))) is not None:
+                    savings = dist - other - 2
+                    if savings >= limit:
+                        total += 1
+
+        return total
+
+    @classmethod
+    def part2(cls, input: str, limit: int = 100) -> int:
+        path = parse_path(input)
+
+        return sum(
+            get_savings(a, b) >= limit
+            for a, b in itertools.combinations(path.items(), 2)
+        )