AoC 2022 Day 12

# aoc202212.py

import pathlib
import sys

from queue import PriorityQueue
from string import ascii_lowercase
from typing import Optional, TypeAlias

# Type aliases
Coordinate: TypeAlias = tuple[int, int]
Grid: TypeAlias = dict[Coordinate, int]
Path: TypeAlias = dict[Coordinate, Optional[Coordinate]]

def parse(puzzle_input: str) -> Grid:
    """Parse input"""
    letter_vals: dict[str, int] = {
        letter: value for (letter, value) in zip(ascii_lowercase, range(1, 27))
    }
    letter_vals["S"] = 0
    letter_vals["E"] = 27
    return {
        (x, y): letter_vals[c]
        for y, line in enumerate(puzzle_input.splitlines())
        for x, c in enumerate(line)
    }


def neighbours(grid: Grid, pos: Coordinate) -> list[Coordinate]:
    """ Return neighbours of a position in a grid """
    x0, y0 = pos
    candidates = [(x0 - 1, y0), (x0 + 1, y0), (x0, y0 - 1), (x0, y0 + 1)]
    return [p for p in candidates if p in grid]


def reconstruct_path(locations: Path,
                     start: Coordinate,
                     goal: Coordinate) -> list[Coordinate]:
    """ Get shortest path """
    here: Coordinate = goal
    path: list[Coordinate] = []
    if goal not in locations:
        return []
    while here != start:
        path.append(here)
        here = locations[here]
    return path


def dijkstra_search(graph: Grid,
                    start: Coordinate,
                    goal: Coordinate,
                    edges: dict[Coordinate, list[Coordinate]]) -> Path:
    """ Find a path using Dijkstra search """
    frontier = PriorityQueue()
    frontier.put(start, 0)
    came_from: Path = {}
    came_from[start] = None
    cost_so_far: dict[Coordinate, int] = {}
    cost_so_far[start] = 0

    while not frontier.empty():
        current: Coordinate = frontier.get()
        if current == goal:
            break
        for edge in edges[current]:
            new_cost: int = cost_so_far[current] + 1
            if not graph[edge] <= graph[current] + 1:
                continue
            if edge not in cost_so_far or new_cost < cost_so_far[edge]:
                cost_so_far[edge] = new_cost
                priority = new_cost
                frontier.put(edge, priority)
                came_from[edge] = current
    return came_from


def part1(grid_dict: Grid) -> int:
    """Solve part 1"""
    start: Coordinate = next(key for key, val in grid_dict.items() if val == 0)
    finish: Coordinate = next(key for key, val in grid_dict.items() if val == 27)
    edges: dict[Coordinate, list[Coordinate]] = {
        coord: neighbours(grid_dict, coord) for coord in grid_dict.keys()
    }

    route: Path = dijkstra_search(grid_dict, start, finish, edges)
    return len(reconstruct_path(route, start, finish))


def part2(grid_dict: Grid) -> int:
    """Solve part 2"""
    finish: Coordinate = next(key for key, val in grid_dict.items() if val == 27)
    lowest_elevations = (key for key, val in grid_dict.items() if val == 1)
    edges: dict[Coordinate, list[Coordinate]] = {
        coord: neighbours(grid_dict, coord) for coord in grid_dict.keys()
    }

    shortest: int = part1(grid_dict)

    for el in lowest_elevations:
        route: Path = dijkstra_search(grid_dict, el, finish, edges)
        length = len(reconstruct_path(route, el, finish))
        if length < shortest and length > 0:
            shortest = length

    return shortest


def solve(puzzle_input: str) -> tuple[int, int]:
    """Solve the puzzle for the given input"""
    data = parse(puzzle_input)
    solution1 = part1(data)  # Correct answer was 394 (with my data)
    solution2 = part2(data)  # Correct answer was 388 (with my data)

    return solution1, solution2


if __name__ == "__main__":
    for path in sys.argv[1:]:
        print(f"{path}:")
        puzzle_input = pathlib.Path(path).read_text().strip()
        solutions = solve(puzzle_input)
        print("\n".join(str(solution) for solution in solutions))