AoC 2024 Day 20 [faster / multiprocessing]

Author: pareronia

src/main/python/AoC2024_20.py

@@ -3,6 +3,8 @@
 # Advent of Code 2024 Day 20
 #
 
+import multiprocessing
+import os
 import sys
 
 from aoc.common import InputData
@@ -39,6 +41,30 @@ def parse_input(self, input_data: InputData) -> Input:
         return CharGrid.from_strings(list(input_data))
 
     def solve(self, grid: CharGrid, cheat_len: int, target: int) -> int:
+        ans = multiprocessing.Manager().dict()
+
+        def count_shortcuts(
+            id: str,
+            cells: list[tuple[int, int]],
+            dist: dict[tuple[int, int], int],
+        ) -> None:
+            count = 0
+            for r, c in cells:
+                for md in range(2, cheat_len + 1):
+                    for dr in range(md + 1):
+                        dc = md - dr
+                        for rr, cc in {
+                            (r + dr, c + dc),
+                            (r + dr, c - dc),
+                            (r - dr, c + dc),
+                            (r - dr, c - dc),
+                        }:
+                            if (rr, cc) not in dist:
+                                continue
+                            if dist[(rr, cc)] - dist[(r, c)] >= target + md:
+                                count += 1
+            ans[id] = count
+
         start = next(
             cell for cell in grid.get_cells() if grid.get_value(cell) == "S"
         )
@@ -52,22 +78,25 @@ def solve(self, grid: CharGrid, cheat_len: int, target: int) -> int:
             ),
         )
         dist = {(k.row, k.col): v for k, v in distances.items()}
-        ans = 0
-        for r, c in dist.keys():
-            for md in range(2, cheat_len + 1):
-                for dr in range(md + 1):
-                    dc = md - dr
-                    for rr, cc in {
-                        (r + dr, c + dc),
-                        (r + dr, c - dc),
-                        (r - dr, c + dc),
-                        (r - dr, c - dc),
-                    }:
-                        if (rr, cc) not in dist:
-                            continue
-                        if dist[(rr, cc)] - dist[(r, c)] >= target + md:
-                            ans += 1
-        return ans
+        if sys.platform.startswith("win"):
+            count_shortcuts("main", [(r, c) for r, c in dist.keys()], dist)
+        else:
+            n = os.process_cpu_count()
+            cells: list[list[tuple[int, int]]] = [[] for _ in range(n)]
+            cnt = 0
+            for r, c in dist.keys():
+                cells[cnt % n].append((r, c))
+                cnt += 1
+            jobs = []
+            for i in range(n):
+                p = multiprocessing.Process(
+                    target=count_shortcuts, args=(str(i), cells[i], dist)
+                )
+                jobs.append(p)
+                p.start()
+            for p in jobs:
+                p.join()
+        return sum(ans.values())
 
     def part_1(self, grid: Input) -> Output1:
         return self.solve(grid, 2, 100)