diff --git a/src/holt59/aoc/2024/day21.py b/src/holt59/aoc/2024/day21.py index 656c478..6e79eee 100644 --- a/src/holt59/aoc/2024/day21.py +++ b/src/holt59/aoc/2024/day21.py @@ -1,121 +1,85 @@ -import heapq -from dataclasses import dataclass -from typing import Any, Iterator, Literal, Sequence, TypeAlias, cast +import itertools +from functools import cache +from typing import Any, Iterator, Literal 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", ">")) +NUM_PAD_P = { + v: (i, j) + for i, r in enumerate(("789", "456", "123", " 0A")) + for j, v in enumerate(r) + if v.strip() +} +MOV_PAD_P = { + v: (i, j) + for i, r in enumerate((" ^A", "")) + for j, v in enumerate(r) + if v.strip() +} -@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) +def path(start: tuple[int, int], end: tuple[int, int], pad: Literal["num", "mov"]): + # a move in the grid is composed of at most two straight line: up/down and + # left/right, since doing some kind of diagonal moves would create long path for + # the robot above (since this involves going back-and-forth to the letter 'A') + # + row_s, col_s = start + row_e, col_e = end - code: str = "" + le, de, ue, re = ( + "<" * max(0, col_s - col_e), + "v" * max(0, row_e - row_s), + "^" * max(0, row_s - row_e), + ">" * max(0, col_e - col_s), + ) + + # when the robot starts or ends on the row/column with the empty cell, there is + # only one way to move + # + if pad == "num" and (row_s, col_e) == (3, 0): + return ue + le + elif pad == "num" and (col_s, row_e) == (0, 3): + return re + de + elif pad == "mov" and col_s == 0: + return re + ue + elif pad == "mov" and col_e == 0: + return de + le + + # otherwise, we need to decide if we want to go up/down first, or left/right, and + # apparently this is the best way to do it... + return le + de + ue + re -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 +@cache +def v_clicks(clicks: str, depth: int) -> int: + if depth == 0: + return len(clicks) - if 0 <= row < len(pad) and 0 <= col < len(pad[row]) and pad[row][col] is not None: - return (row, col) + n_clicks = 0 + at = "A" + for _, group in itertools.groupby(clicks): + group = list(group) + n_clicks += v_clicks( + path(MOV_PAD_P[at], MOV_PAD_P[group[0]], "mov") + "A" * len(group), + depth - 1, + ) + at = group[0] - return None + return n_clicks -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), +def path_length(code: str, depth: int): + return sum( + v_clicks(path(NUM_PAD_P[start], NUM_PAD_P[end], "num") + "A", depth) + for start, end in zip("A" + code[:-1], code, strict=True) ) class Solver(BaseSolver): - 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() + path_length(code, 2) * int(code[:-1], 10) for code in input.splitlines() + ) + yield sum( + path_length(code, 25) * int(code[:-1], 10) for code in input.splitlines() )