awful day 21 solution part b python

AOC/python/solutions/day_21.py

@@ -1,5 +1,6 @@
 from functools import cache
-
+from itertools import groupby
+from models.aoc_solution import Dataset
 from helpers.grid_helpers import find_char
 from models.aoc_solution import AOCSolution
 
@@ -9,7 +10,7 @@
 class Day21(AOCSolution):
     EXPECTED = {
         "part_one": {"sample": 126384, "data": 278748},
-        "part_two": {"sample": 0, "data": 0},
+        "part_two": {"sample": 0, "data": 337744744231414},
     }
 
     def __post_init__(self) -> None:
@@ -26,10 +27,12 @@ def __post_init__(self) -> None:
         ]
         self.key_start = find_char(self.keypad, "A")
         self.dir_start = find_char(self.directions, "A")
-        self.numpad_paths.cache_clear()
+        self.paths.cache_clear()
+        self.code_paths.cache_clear()
+        self.scale_path.cache_clear()
 
     @cache
-    def numpad_paths(self, start: str, end: str, is_numpad: bool) -> set[str]:
+    def paths(self, start: str, end: str, is_numpad: bool) -> set[str]:
         """Generate all paths between start and end avoiding None tiles"""
         grid = self.keypad if is_numpad else self.directions
         x, y = find_char(grid, start)
@@ -38,44 +41,65 @@ def numpad_paths(self, start: str, end: str, is_numpad: bool) -> set[str]:
         paths: set[str] = set()
         path = f"{'<>'[dx > 0] * abs(dx)}{'^v'[dy > 0] * abs(dy)}"
         if grid[y][x + dx] is not None:
-            paths.add(f"{path}A")
+            paths.add(path)
         if grid[y + dy][x] is not None:
-            paths.add(f"{path[::-1]}A")
+            paths.add(path[::-1])
         return paths
 
+    @cache
     def code_paths(self, code: str, is_numpad: bool) -> set[str]:
         """Generate all paths for a given code"""
         if len(code) < 2:
             return {""}
-        starts = self.numpad_paths(code[0], code[1], is_numpad)
+        starts = self.paths(code[0], code[1], is_numpad)
         return {
-            start + path
+            f"{start}A{path}"
             for start in starts
             for path in self.code_paths(code[1:], is_numpad)
         }
 
-    def generate_parent_paths(self, code: str) -> set[str]:
+    def generate_parent_paths(self, code: str, robot_count: int) -> set[str]:
         """Generate all paths for a given code"""
-        paths = self.code_paths(f"A{code}", True)
-        parent_paths = set()
-        for _ in range(2):
+        parent_paths = paths = self.code_paths(f"A{code}", True)
+        for _ in range(robot_count):
             parent_paths = set()
             for path in paths:
                 parent_paths.update(self.code_paths(f"A{path}", False))
             paths = parent_paths
         return parent_paths
 
+    def get_shortest_path(self, code: str, robot_count: int) -> str:
+        paths = self.generate_parent_paths(code, robot_count)
+        return min(paths, key=len)
+
+    @cache
+    def scale_path(self, path: str, robot_count: int) -> int:
+        """What is the length of the path after passing through the given number of robots"""
+        scaled = 0
+        prev = "A"
+        for char, group in groupby(path):
+            length = len(list(group))
+            shortest = min(self.paths(prev, char, False), key=len)
+            if robot_count == 1:
+                scaled += len(shortest)
+                scaled += length
+            else:
+                scaled += self.scale_path(f"{shortest}{'A' * length}", robot_count - 1)
+            prev = char
+        return scaled
+
     def part_one(self) -> int:
-        total = 0
-        for code in self.codes:
-            paths = self.generate_parent_paths(code)
-            shortest = min(len(path) for path in paths)
-            numeric = int(code[:-1])
-            total += numeric * shortest
-        return total
+        return sum(
+            len(self.get_shortest_path(code, 2)) * int(code[:-1]) for code in self.codes
+        )
 
     def part_two(self) -> int:
-        return 0
+        """Turns out this is not deterministic, so run it 1000 times and take the best one
+        TODO: fix this"""
+        return sum(
+            self.scale_path(self.get_shortest_path(code, 0), 25) * int(code[:-1])
+            for code in self.codes
+        )
 
 
 if __name__ == "__main__":