from __future__ import annotations import heapq import os import sys from collections import defaultdict from dataclasses import dataclass from typing import Literal, TypeAlias VERBOSE = os.getenv("AOC_VERBOSE") == "True" Direction: TypeAlias = Literal[">", "<", "^", "v"] @dataclass(frozen=True, order=True) class Label: row: int col: int direction: Direction parent: Label | None = None # mappings from direction to row shift / col shift / opposite direction MAPPINGS: dict[Direction, tuple[int, int, Direction]] = { ">": (0, +1, "<"), "<": (0, -1, ">"), "v": (+1, 0, "^"), "^": (-1, 0, "v"), } def print_shortest_path( grid: list[list[int]], target: tuple[int, int], per_cell: dict[tuple[int, int], list[tuple[Label, int]]], ): assert len(per_cell[target]) == 1 label = per_cell[target][0][0] path: list[Label] = [] while True: path.insert(0, label) if label.parent is None: break label = label.parent p_grid = [[str(c) for c in r] for r in grid] for i in range(len(grid)): for j in range(len(grid[0])): if per_cell[i, j]: p_grid[i][j] = f"\033[94m{grid[i][j]}\033[0m" prev_label = path[0] for label in path[1:]: for r in range( min(prev_label.row, label.row), max(prev_label.row, label.row) + 1 ): for c in range( min(prev_label.col, label.col), max(prev_label.col, label.col) + 1, ): if (r, c) != (prev_label.row, prev_label.col): p_grid[r][c] = f"\033[93m{grid[r][c]}\033[0m" p_grid[label.row][label.col] = f"\033[91m{grid[label.row][label.col]}\033[0m" prev_label = label p_grid[0][0] = f"\033[92m{grid[0][0]}\033[0m" print("\n".join("".join(row) for row in p_grid)) def shortest_many_paths(grid: list[list[int]]) -> dict[tuple[int, int], int]: n_rows, n_cols = len(grid), len(grid[0]) visited: dict[tuple[int, int], tuple[Label, int]] = {} queue: list[tuple[int, Label]] = [ (0, Label(row=n_rows - 1, col=n_cols - 1, direction="^")) ] while queue and len(visited) != n_rows * n_cols: distance, label = heapq.heappop(queue) if (label.row, label.col) in visited: continue visited[label.row, label.col] = (label, distance) for direction, (c_row, c_col, i_direction) in MAPPINGS.items(): if label.direction == i_direction: continue else: row, col = (label.row + c_row, label.col + c_col) # exclude labels outside the grid or with too many moves in the same # direction if row not in range(0, n_rows) or col not in range(0, n_cols): continue heapq.heappush( queue, ( distance + sum( grid[r][c] for r in range(min(row, label.row), max(row, label.row) + 1) for c in range(min(col, label.col), max(col, label.col) + 1) ) - grid[row][col], Label( row=row, col=col, direction=direction, parent=label, ), ), ) return {(r, c): visited[r, c][1] for r in range(n_rows) for c in range(n_cols)} def shortest_path( grid: list[list[int]], min_straight: int, max_straight: int, lower_bounds: dict[tuple[int, int], int], ) -> int: n_rows, n_cols = len(grid), len(grid[0]) target = (len(grid) - 1, len(grid[0]) - 1) # for each tuple (row, col, direction, count), the associated label when visited visited: dict[tuple[int, int, Direction], Label] = {} # list of all visited labels for a cell (with associated distance) per_cell: dict[tuple[int, int], list[tuple[Label, int]]] = defaultdict(list) # need to add two start labels, otherwise one of the two possible direction will # not be possible queue: list[tuple[int, int, Label]] = [ (lower_bounds[0, 0], 0, Label(row=0, col=0, direction="^")), (lower_bounds[0, 0], 0, Label(row=0, col=0, direction="<")), ] while queue: _, distance, label = heapq.heappop(queue) if (label.row, label.col, label.direction) in visited: continue visited[label.row, label.col, label.direction] = label per_cell[label.row, label.col].append((label, distance)) if (label.row, label.col) == target: break for direction, (c_row, c_col, i_direction) in MAPPINGS.items(): # cannot move in the opposite direction if label.direction == i_direction or label.direction == direction: continue distance_to = distance for amount in range(1, max_straight + 1): row, col = ( label.row + amount * c_row, label.col + amount * c_col, ) # exclude labels outside the grid or with too many moves in the same # direction if not (0 <= row < n_rows and 0 <= col < n_cols): break distance_to += grid[row][col] if amount < min_straight: continue heapq.heappush( queue, ( distance_to + lower_bounds[row, col], distance_to, Label( row=row, col=col, direction=direction, parent=label, ), ), ) if VERBOSE: print_shortest_path(grid, target, per_cell) return per_cell[target][0][1] data = [[int(c) for c in r] for r in sys.stdin.read().splitlines()] estimates = shortest_many_paths(data) # part 1 answer_1 = shortest_path(data, 1, 3, lower_bounds=estimates) print(f"answer 1 is {answer_1}") # part 2 answer_2 = shortest_path(data, 4, 10, lower_bounds=estimates) print(f"answer 2 is {answer_2}")