Implement 2024 day 21 part 1

2024/src/aoc/days/day21.py

@@ -0,0 +1,132 @@
+import functools
+
+from . import SeparateRunner
+
+NUMPAD = {
+    "A": (3, 2),
+    "0": (3, 1),
+    "1": (2, 0),
+    "2": (2, 1),
+    "3": (2, 2),
+    "4": (1, 0),
+    "5": (1, 1),
+    "6": (1, 2),
+    "7": (0, 0),
+    "8": (0, 1),
+    "9": (0, 2),
+}
+
+DIRPAD = {
+    "A": (0, 2),
+    "^": (0, 1),
+    "<": (1, 0),
+    "v": (1, 1),
+    ">": (1, 2),
+}
+
+
+@functools.cache
+def shortest_numpad(from_: str, to: str) -> list[str]:
+    inverse = set(NUMPAD.values())
+    ay, ax = NUMPAD[from_]
+    by, bx = NUMPAD[to]
+
+    dx, dy = bx - ax, by - ay
+
+    sx = "<" if dx < 0 else ">"
+    sy = "^" if dy < 0 else "v"
+
+    if dx > 0 and (by, ax) in inverse or (ay, bx) not in inverse:
+        return abs(dy) * sy + abs(dx) * sx + "A"
+    else:
+        return abs(dx) * sx + abs(dy) * sy + "A"
+
+
+@functools.cache
+def shortest_dirpad(from_: str, to: str) -> str:
+    inverse = set(DIRPAD.values())
+    ay, ax = DIRPAD[from_]
+    by, bx = DIRPAD[to]
+
+    dx, dy = bx - ax, by - ay
+    sx = "<" if dx < 0 else ">"
+    sy = "^" if dy < 0 else "v"
+
+    if dx > 0 and (by, ax) in inverse or (ay, bx) not in inverse:
+        return abs(dy) * sy + abs(dx) * sx + "A"
+    else:
+        return abs(dx) * sx + abs(dy) * sy + "A"
+
+
+def encode_shortest_numpad(code: str) -> str:
+    pos = "A"
+
+    res = ""
+
+    for c in code:
+        res += shortest_numpad(pos, c)
+        # print(c, res)
+        pos = c
+
+    return res
+
+
+def encode_shortest_dirpad(code: str) -> str:
+    pos = "A"
+
+    res = ""
+
+    for c in code:
+        if pos != c:
+            res += shortest_dirpad(pos, c)
+        else:
+            res += "A"
+        pos = c
+
+    return res
+
+
+def decode(code: str, pad: dict[str, tuple[int, int]]) -> str:
+    result = ""
+    inverse = {v: k for k, v in pad.items()}
+
+    y, x = pad["A"]
+
+    for i, c in enumerate(code):
+        match c:
+            case "A":
+                result += inverse[y, x]
+            case "^":
+                y -= 1
+            case "v":
+                y += 1
+            case "<":
+                x -= 1
+            case ">":
+                x += 1
+
+        if (y, x) not in inverse:
+            raise ValueError(
+                f"""Moved off the board {x, y}, after processing {c}.
+                Path so far: {result} (from {code[:i]})"""
+            )
+
+    return result
+
+
+class DayRunner(SeparateRunner):
+    @classmethod
+    def part1(cls, input: str) -> int:
+        result = 0
+        for code in input.strip().split("\n"):
+            numpad = encode_shortest_numpad(code)
+            robot1 = encode_shortest_dirpad(numpad)
+            robot2 = encode_shortest_dirpad(robot1)
+
+            result += int(code[:-1]) * len(robot2)
+
+        return result
+
+    @classmethod
+    def part2(cls, input: str) -> int:
+        pass