From ae5527b72dde549e456208c2be827a64ae61771e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mika=C3=ABl=20Capelle?= <capelle.mikael@gmail.com>
Date: Sat, 21 Dec 2024 13:13:31 +0100
Subject: [PATCH] 2024 day 21 part 1.

---
 src/holt59/aoc/2024/day21.py                | 118 +++++++++++++++++++-
 src/holt59/aoc/inputs/holt59/2024/day21.txt |   5 +
 src/holt59/aoc/inputs/tests/2024/day21.txt  |   5 +
 3 files changed, 126 insertions(+), 2 deletions(-)

diff --git a/src/holt59/aoc/2024/day21.py b/src/holt59/aoc/2024/day21.py
index 07e201e..656c478 100644
--- a/src/holt59/aoc/2024/day21.py
+++ b/src/holt59/aoc/2024/day21.py
@@ -1,7 +1,121 @@
-from typing import Any, Iterator
+import heapq
+from dataclasses import dataclass
+from typing import Any, Iterator, Literal, Sequence, TypeAlias, cast
 
 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", ">"))
+
+
+@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)
+
+    code: str = ""
+
+
+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
+
+    if 0 <= row < len(pad) and 0 <= col < len(pad[row]) and pad[row][col] is not None:
+        return (row, col)
+
+    return None
+
+
+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),
+    )
+
 
 class Solver(BaseSolver):
-    def solve(self, input: str) -> Iterator[Any]: ...
+    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()
+        )
diff --git a/src/holt59/aoc/inputs/holt59/2024/day21.txt b/src/holt59/aoc/inputs/holt59/2024/day21.txt
index e69de29..bb898ec 100644
--- a/src/holt59/aoc/inputs/holt59/2024/day21.txt
+++ b/src/holt59/aoc/inputs/holt59/2024/day21.txt
@@ -0,0 +1,5 @@
+129A
+540A
+789A
+596A
+582A
diff --git a/src/holt59/aoc/inputs/tests/2024/day21.txt b/src/holt59/aoc/inputs/tests/2024/day21.txt
index e69de29..4cf0c29 100644
--- a/src/holt59/aoc/inputs/tests/2024/day21.txt
+++ b/src/holt59/aoc/inputs/tests/2024/day21.txt
@@ -0,0 +1,5 @@
+029A
+980A
+179A
+456A
+379A