diff --git a/2024/d21/ex1/ex1.py b/2024/d21/ex1/ex1.py new file mode 100755 index 0000000..381c2c0 --- /dev/null +++ b/2024/d21/ex1/ex1.py @@ -0,0 +1,108 @@ +#!/usr/bin/env python + +import enum +import functools +import itertools +import sys +from typing import Literal, NamedTuple + + +class Point(NamedTuple): + x: int + y: int + + +class Direction(enum.StrEnum): + UP = "^" + DOWN = "v" + LEFT = "<" + RIGHT = ">" + + +Instruction = Direction | Literal["A"] + +PAD = { + "7": Point(-3, -2), + "8": Point(-3, -1), + "9": Point(-3, 0), + "4": Point(-2, -2), + "5": Point(-2, -1), + "6": Point(-2, 0), + "1": Point(-1, -2), + "2": Point(-1, -1), + "3": Point(-1, 0), + "0": Point(0, -1), + "A": Point(0, 0), +} + +ARROWS = { + "^": Point(0, -1), + "A": Point(0, 0), + "<": Point(1, -2), + "v": Point(1, -1), + ">": Point(1, 0), +} + +# Needs to be hash-able +InstructionList = tuple[Instruction, ...] + + +def solve(input: str) -> int: + def button_paths( + start_button: str, + end_button: str, + buttons: dict[str, Point], + ) -> set[InstructionList]: + start, end = buttons[start_button], buttons[end_button] + sequences: set[InstructionList] = set() + + dx, dy = (end.x - start.x), (end.y - start.y) + + a_button: list[Instruction] = ["A"] # Work around MyPy limitation + move_x = [Direction.UP if dx < 0 else Direction.DOWN] * abs(dx) + move_y = [Direction.LEFT if dy < 0 else Direction.RIGHT] * abs(dy) + + # Avoid moving over the gap + if Point(end.x, start.y) in buttons.values(): + sequences.add(tuple(a_button + move_x + move_y + a_button)) + if Point(start.x, end.y) in buttons.values(): + sequences.add(tuple(a_button + move_y + move_x + a_button)) + + return sequences + + @functools.cache + def path_cost(path: InstructionList, depth: int) -> int: + # Have we reached the actual keypad robot + if depth == 0: + # We start on 'A' so don't count it + return len(path) - 1 + # Otherwise, intermediate robot must use arrow pad + cost = sum( + min(path_cost(path, depth - 1) for path in button_paths(start, end, ARROWS)) + for start, end in itertools.pairwise(path) + ) + return cost + + def code_cost(code: str, depth: int) -> int: + # We start on 'A' + code = "A" + code + cost = sum( + min(path_cost(path, depth) for path in button_paths(start, end, PAD)) + for start, end in itertools.pairwise(code) + ) + return cost + + def complexity(code: str, seq_len: int) -> int: + return int(code.replace("A", "")) * seq_len + + codes = input.splitlines() + return sum(complexity(code, code_cost(code, 2)) for code in codes) + + +def main() -> None: + input = sys.stdin.read() + print(solve(input)) + + +if __name__ == "__main__": + main()