Merge pull request #92 from vil02/21_solution

feat: add `adv_2024_21`

solutions/adv_2024_21.py

@@ -0,0 +1,147 @@
+import typing
+import functools
+import itertools
+
+
+Keypad = dict[str, dict[str, str]]
+
+
+_NUMERIC_KEYPAD = {
+    "A": {"<": "0", "^": "3"},
+    "0": {"^": "2", ">": "A"},
+    "1": {"^": "4", ">": "2"},
+    "2": {"<": "1", ">": "3", "^": "5", "v": "0"},
+    "3": {"<": "2", "^": "6", "v": "A"},
+    "4": {"^": "7", "v": "1", ">": "5"},
+    "5": {"<": "4", ">": "6", "^": "8", "v": "2"},
+    "6": {"<": "5", "^": "9", "v": "3"},
+    "7": {"v": "4", ">": "8"},
+    "8": {"<": "7", "v": "5", ">": "9"},
+    "9": {"<": "8", "v": "6"},
+}
+
+assert all(len(_NUMERIC_KEYPAD[_]) == 4 for _ in ["2", "5"])
+assert all(len(_NUMERIC_KEYPAD[_]) == 3 for _ in ["4", "6", "8"])
+assert all(len(_NUMERIC_KEYPAD[_]) == 2 for _ in ["1", "0", "A", "9", "7"])
+
+
+_DIRECTIONAL_KEYPAD = {
+    "^": {">": "A", "v": "v"},
+    "A": {"<": "^", "v": ">"},
+    "<": {">": "v"},
+    "v": {"^": "^", ">": ">", "<": "<"},
+    ">": {"^": "A", "<": "v"},
+}
+
+_NEGATE_DIR = {
+    "<": ">",
+    ">": "<",
+    "v": "^",
+    "^": "v",
+}
+
+
+def _does_make_sense(path: str) -> bool:
+    moves = set(_ for _ in path)
+    return all(_NEGATE_DIR[_] not in moves for _ in moves)
+
+
+def _find_all_sequences(keypad: Keypad, start: str, end: str) -> set[str]:
+    active = [(start, "")]
+    visited = set()
+    res: set[str] = set()
+    while active:
+        cur_key, cur_path = active.pop()
+        if cur_key == end:
+            res.add(cur_path)
+            continue
+        assert tuple(cur_path) not in visited
+        visited.add(tuple(cur_path))
+        for new_dir, new_key in keypad[cur_key].items():
+            new_path = cur_path + new_dir
+            if _does_make_sense(new_path):
+                active.append((new_key, new_path))
+    assert len({len(_) for _ in res}) == 1
+    return res
+
+
+@functools.cache
+def numeric_keypad_sequences(start: str, end: str) -> set[str]:
+    return _find_all_sequences(_NUMERIC_KEYPAD, start, end)
+
+
+@functools.cache
+def directional_keypad_sequences(start: str, end: str) -> set[str]:
+    return _find_all_sequences(_DIRECTIONAL_KEYPAD, start, end)
+
+
+def _combine_sequences(seq_a: set[str], seq_b: set[str]) -> set[str]:
+    return {_a + _b for _a, _b in itertools.product(seq_a, seq_b)}
+
+
+def _append_with_a(seq: set[str]) -> set[str]:
+    return {_ + "A" for _ in seq}
+
+
+def _get_keypad_press_sequence(
+    keypad_sequences: typing.Callable[[str, str], set[str]]
+) -> typing.Callable[[str], set[str]]:
+    def _keypad_press_sequence(keys: str) -> set[str]:
+        res = {""}
+        prev_key = "A"
+        for _ in keys:
+            res = _append_with_a(_combine_sequences(res, keypad_sequences(prev_key, _)))
+            prev_key = _
+        return res
+
+    return _keypad_press_sequence
+
+
+numeric_keypad_press_sequence = _get_keypad_press_sequence(numeric_keypad_sequences)
+directional_keypad_press_sequence = _get_keypad_press_sequence(
+    directional_keypad_sequences
+)
+
+
+def _subproblems(code: str) -> list[str]:
+    assert code.endswith("A")
+    pieces = code[:-1].split("A")
+    return [_ + "A" for _ in pieces]
+
+
+@functools.cache
+def _shortest(code: str, iterations: int) -> int:
+    if iterations == 0:
+        return len(code)
+    return min(
+        sum(_shortest(_, iterations - 1) for _ in _subproblems(cur_seq))
+        for cur_seq in directional_keypad_press_sequence(code)
+    )
+
+
+def shortest_full_sequence_press(code: str, iter_size: int) -> int:
+    res = numeric_keypad_press_sequence(code)
+    return min(_shortest(_, iter_size) for _ in res)
+
+
+def numeric_part(code: str) -> int:
+    return int(code.replace("A", ""))
+
+
+def _code_complexity(code: str, iter_size: int) -> int:
+    return shortest_full_sequence_press(code, iter_size) * numeric_part(code)
+
+
+def _parse_input(in_str: str) -> list[str]:
+    return in_str.splitlines()
+
+
+def _get_solve(iter_size: int) -> typing.Callable[[str], int]:
+    def _solve(in_str: str) -> int:
+        return sum(_code_complexity(_, iter_size) for _ in _parse_input(in_str))
+
+    return _solve
+
+
+solve_a = _get_solve(2)
+solve_b = _get_solve(25)

tests/test_adv_2024_21.py

@@ -0,0 +1,99 @@
+import pytest
+
+import solutions.adv_2024_21 as sol
+from . import test_utils as tu
+
+
+@pytest.mark.parametrize(
+    ("start", "end", "expected"),
+    [
+        ("A", "A", {""}),
+        ("A", "3", {"^"}),
+        ("5", "9", {">^", "^>"}),
+        ("3", "4", {"<<^", "<^<", "^<<"}),
+        ("9", "1", {"<<vv", "vv<<", "<v<v", "v<v<", "<vv<", "v<<v"}),
+    ],
+)
+def test_numeric_keypad_sequences(start: str, end: str, expected: set[str]) -> None:
+    assert sol.numeric_keypad_sequences(start, end) == expected
+
+
+@pytest.mark.parametrize(
+    ("start", "end", "expected"),
+    [
+        ("A", "A", {""}),
+        ("<", "v", {">"}),
+        ("<", "^", {">^"}),
+        ("<", "A", {">>^", ">^>"}),
+        ("A", "v", {"<v", "v<"}),
+        ("A", "<", {"<v<", "v<<"}),
+    ],
+)
+def test_directional_keypad_sequences(start: str, end: str, expected: set[str]) -> None:
+    assert sol.directional_keypad_sequences(start, end) == expected
+
+
+@pytest.mark.parametrize(
+    ("keys", "expected"),
+    [
+        ("029A", {"<A^A>^^AvvvA", "<A^A^>^AvvvA", "<A^A^^>AvvvA"}),
+    ],
+)
+def test_numeric_keypad_press_sequence(keys: str, expected: set[str]) -> None:
+    assert sol.numeric_keypad_press_sequence(keys) == expected
+
+
+@pytest.mark.parametrize(
+    ("code", "iter_size", "expected"),
+    [
+        (
+            "029A",
+            2,
+            len("<vA<AA>>^AvAA<^A>A<v<A>>^AvA^A<vA>^A<v<A>^A>AAvA^A<v<A>A>^AAAvA<^A>A"),
+        ),
+        (
+            "980A",
+            2,
+            len("<v<A>>^AAAvA^A<vA<AA>>^AvAA<^A>A<v<A>A>^AAAvA<^A>A<vA>^A<A>A"),
+        ),
+        (
+            "179A",
+            2,
+            len("<v<A>>^A<vA<A>>^AAvAA<^A>A<v<A>>^AAvA^A<vA>^AA<A>A<v<A>A>^AAAvA<^A>A"),
+        ),
+        (
+            "456A",
+            2,
+            len("<v<A>>^AA<vA<A>>^AAvAA<^A>A<vA>^A<A>A<vA>^A<A>A<v<A>A>^AAvA<^A>A"),
+        ),
+        (
+            "379A",
+            2,
+            len("<v<A>>^AvA^A<vA<AA>>^AAvA<^A>AAvA^A<vA>^AA<A>A<v<A>A>^AAAvA<^A>A"),
+        ),
+    ],
+)
+def test_shortest_full_sequence_press(code: str, iter_size: int, expected: int) -> None:
+    assert sol.shortest_full_sequence_press(code, iter_size) == expected
+
+
+@pytest.mark.parametrize(
+    ("code", "expected"),
+    [
+        ("029A", 29),
+        ("980A", 980),
+        ("179A", 179),
+        ("456A", 456),
+        ("379A", 379),
+    ],
+)
+def test_numeric_part(code: str, expected: int) -> None:
+    assert sol.numeric_part(code) == expected
+
+
+_INPUTS = tu.get_inputs(21, {"small", "p"})
+
+test_solve_a, test_solve_b = _INPUTS.get_tests(
+    (sol.solve_a, sol.solve_b),
+    {"small": (126384, 154115708116294), "p": (278568, 341460772681012)},
+)

tests/test_input_data/data_adv_2024_21_p.txt

@@ -0,0 +1,5 @@
+540A
+839A
+682A
+826A
+974A

tests/test_input_data/data_adv_2024_21_small.txt

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