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

2024/tests/samples/21.txt

@@ -0,0 +1,5 @@
+029A
+980A
+179A
+456A
+379A

2024/tests/test_day21.py

@@ -0,0 +1,51 @@
+import pytest
+
+from aoc.days.day21 import (
+    DayRunner,
+    encode_shortest_dirpad,
+    encode_shortest_numpad,
+)
+
+from . import get_data
+
+
+def test_encode_shortest_numpad() -> None:
+    assert encode_shortest_numpad("029A") in (
+        "<A^A>^^AvvvA",
+        "<A^A^>^AvvvA",
+        "<A^A^^>AvvvA",
+    )
+
+
+def test_encode_shortest_dirpad() -> None:
+    numpad_encoded = encode_shortest_numpad("029A")
+    assert len(encode_shortest_dirpad(numpad_encoded)) == len(
+        "v<<A>>^A<A>AvA<^AA>A<vAAA>^A"
+    )
+
+
+@pytest.mark.parametrize(
+    "code,answer",
+    [
+        (
+            "029A",
+            "<vA<AA>>^AvAA<^A>A<v<A>>^AvA^A<vA>^A<v<A>^A>AAvA^A<v<A>A>^AAAvA<^A>A",
+        ),
+        ("980A", "<v<A>>^AAAvA^A<vA<AA>>^AvAA<^A>A<v<A>A>^AAAvA<^A>A<vA>^A<A>A"),
+        (
+            "179A",
+            "<v<A>>^A<vA<A>>^AAvAA<^A>A<v<A>>^AAvA^A<vA>^AA<A>A<v<A>A>^AAAvA<^A>A",
+        ),
+        ("456A", "<v<A>>^AA<vA<A>>^AAvAA<^A>A<vA>^A<A>A<vA>^A<A>A<v<A>A>^AAvA<^A>A"),
+        ("379A", "<v<A>>^AvA^A<vA<AA>>^AAvA<^A>AAvA^A<vA>^AA<A>A<v<A>A>^AAAvA<^A>A"),
+    ],
+)
+def test_encode_shortest_dirpad_twice(code: str, answer: str) -> None:
+    numpad_encoded = encode_shortest_numpad(code)
+    robot1 = encode_shortest_dirpad(numpad_encoded)
+    robot2 = encode_shortest_dirpad(robot1)
+    assert len(robot2) == len(answer)
+
+
+def test_sample_part1() -> None:
+    assert DayRunner.part1(get_data(21)) == 126384