2024 day 21 part 2.

src/holt59/aoc/2024/day21.py

@@ -1,121 +1,85 @@
-import heapq
-from dataclasses import dataclass
-from typing import Any, Iterator, Literal, Sequence, TypeAlias, cast
+import itertools
+from functools import cache
+from typing import Any, Iterator, Literal
 
 from ..base import BaseSolver
 
-Action: TypeAlias = Literal[">", "<", "v", "^", "A"]
-
-NUM_PAD = ((7, 8, 9), (4, 5, 6), (1, 2, 3), (None, 0, "A"))
-MOV_PAD: tuple[tuple[Action | None, ...], ...] = ((None, "^", "A"), ("<", "v", ">"))
+NUM_PAD_P = {
+    v: (i, j)
+    for i, r in enumerate(("789", "456", "123", " 0A"))
+    for j, v in enumerate(r)
+    if v.strip()
+}
+MOV_PAD_P = {
+    v: (i, j)
+    for i, r in enumerate((" ^A", "<v>"))
+    for j, v in enumerate(r)
+    if v.strip()
+}
 
 
-@dataclass(frozen=True, order=True)
-class Node:
-    robot_1: tuple[int, int] = (0, 2)
-    robot_2: tuple[int, int] = (0, 2)
-    robot_3: tuple[int, int] = (3, 2)
+def path(start: tuple[int, int], end: tuple[int, int], pad: Literal["num", "mov"]):
+    # a move in the grid is composed of at most two straight line: up/down and
+    # left/right, since doing some kind of diagonal moves would create long path for
+    # the robot above (since this involves going back-and-forth to the letter 'A')
+    #
+    row_s, col_s = start
+    row_e, col_e = end
 
-    code: str = ""
+    le, de, ue, re = (
+        "<" * max(0, col_s - col_e),
+        "v" * max(0, row_e - row_s),
+        "^" * max(0, row_s - row_e),
+        ">" * max(0, col_e - col_s),
+    )
+
+    # when the robot starts or ends on the row/column with the empty cell, there is
+    # only one way to move
+    #
+    if pad == "num" and (row_s, col_e) == (3, 0):
+        return ue + le
+    elif pad == "num" and (col_s, row_e) == (0, 3):
+        return re + de
+    elif pad == "mov" and col_s == 0:
+        return re + ue
+    elif pad == "mov" and col_e == 0:
+        return de + le
+
+    # otherwise, we need to decide if we want to go up/down first, or left/right, and
+    # apparently this is the best way to do it...
+    return le + de + ue + re
 
 
-def apply_action(
-    robot: tuple[int, int],
-    action: Action,
-    pad: tuple[tuple[int | str | None, ...], ...],
-):
-    d_row, d_col = {"^": (-1, 0), "v": (1, 0), ">": (0, 1), "<": (0, -1)}[action]
-    row, col = robot[0] + d_row, robot[1] + d_col
+@cache
+def v_clicks(clicks: str, depth: int) -> int:
+    if depth == 0:
+        return len(clicks)
 
-    if 0 <= row < len(pad) and 0 <= col < len(pad[row]) and pad[row][col] is not None:
-        return (row, col)
+    n_clicks = 0
+    at = "A"
+    for _, group in itertools.groupby(clicks):
+        group = list(group)
+        n_clicks += v_clicks(
+            path(MOV_PAD_P[at], MOV_PAD_P[group[0]], "mov") + "A" * len(group),
+            depth - 1,
+        )
+        at = group[0]
 
-    return None
+    return n_clicks
 
 
-def create_node(node: Node, action: Action) -> Node | None:
-    # main pad moves -> move first robot
-    if action != "A":
-        robot = apply_action(node.robot_1, action, MOV_PAD)
-        if robot is not None:
-            return Node(
-                robot_1=robot,
-                robot_2=node.robot_2,
-                robot_3=node.robot_3,
-                code=node.code,
-            )
-
-        return None
-
-    # activate pad 1 -> action on robot 1
-    robot_1_action = MOV_PAD[node.robot_1[0]][node.robot_1[1]]
-    assert robot_1_action is not None
-
-    if robot_1_action != "A":
-        robot2 = apply_action(node.robot_2, robot_1_action, MOV_PAD)
-        if robot2 is not None:
-            return Node(
-                robot_1=node.robot_1,
-                robot_2=robot2,
-                robot_3=node.robot_3,
-                code=node.code,
-            )
-        return None
-
-    # activate pad 2 -> action on robot 2
-    robot_2_action = MOV_PAD[node.robot_2[0]][node.robot_2[1]]
-    assert robot_2_action is not None
-
-    if robot_2_action != "A":
-        robot3 = apply_action(node.robot_3, robot_2_action, NUM_PAD)
-        if robot3 is not None:
-            return Node(
-                robot_1=node.robot_1,
-                robot_2=node.robot_2,
-                robot_3=robot3,
-                code=node.code,
-            )
-        return None
-
-    value = NUM_PAD[node.robot_3[0]][node.robot_3[1]]
-    assert value is not None
-    return Node(
-        robot_1=node.robot_1,
-        robot_2=node.robot_2,
-        robot_3=node.robot_3,
-        code=node.code + str(value),
+def path_length(code: str, depth: int):
+    return sum(
+        v_clicks(path(NUM_PAD_P[start], NUM_PAD_P[end], "num") + "A", depth)
+        for start, end in zip("A" + code[:-1], code, strict=True)
     )
 
 
 class Solver(BaseSolver):
-    def dijkstra_for_code(self, target: str):
-        queue: list[tuple[float, Node, tuple[str, ...]]] = [(0, Node(), ())]
-        preds: dict[Node, tuple[str, ...]] = {}
-
-        while queue:
-            dis, node, path = heapq.heappop(queue)
-
-            if not target.startswith(node.code):
-                continue
-
-            if node in preds:
-                continue
-
-            preds[node] = path
-
-            if node.code == target:
-                self.logger.info(f"found [{target}]: {''.join(path)} ({len(path)})")
-                return path
-
-            for action in cast(Sequence[Action], "A^v<>"):
-                node_2 = create_node(node, action)
-                if node_2:
-                    heapq.heappush(queue, (dis + 1, node_2, path + (action,)))
-
-        return None
-
     def solve(self, input: str) -> Iterator[Any]:
         yield sum(
-            len(self.dijkstra_for_code(code) or ()) * int(code[:-1], 10)
-            for code in input.splitlines()
+            path_length(code, 2) * int(code[:-1], 10) for code in input.splitlines()
+        )
+        yield sum(
+            path_length(code, 25) * int(code[:-1], 10) for code in input.splitlines()
         )