AoC 21-15

from typing import List, Tuple
import heapq
import numpy as np

file = "inputs/2021/15/data.txt"
with open(file, "r") as f:
    grid = f.read().split("\n")
    grid = np.array([np.array(list(map(int, line))) for line in grid])

GRID_SHAPE = (len(grid), len(grid[0]))

def get_neighbors(node: Tuple[int, int], grid_shape: Tuple[int, int]) -> List[Tuple[int, int]]:
    row = node[0]
    col = node[1]

    return [
        (row + row_delta, col + col_delta)
        for row_delta in range(-1, 2)
        for col_delta in range(-1, 2)
        if ((row_delta == 0) ^ (col_delta == 0))
        and 0 <= row + row_delta < grid_shape[0]
        and 0 <= col + col_delta < grid_shape[1]
    ]

def manhattan_heuristic(a: Tuple[int, int], b: Tuple[int, int]) -> int:
    return abs(b[0] - a[0]) + abs(b[1] - a[1])

def a_star(start: Tuple[int, int], target: Tuple[int, int], grid: List[List[int]]):
    unexplored_nodes = []
    heapq.heappush(unexplored_nodes, (0, start))

    current_node = None
    prev = {start: None}
    g_scores = {start: 0}

    while current_node != target:
        current_node = heapq.heappop(unexplored_nodes)[1]
        for neighbor in get_neighbors(current_node, grid.shape):
            new_g_score = g_scores[current_node] + grid[neighbor[0]][neighbor[1]]
            if new_g_score < g_scores.get(neighbor, 1e10):
                # Better path to neighbor found
                prev[neighbor] = current_node
                g_scores[neighbor] = new_g_score
                heapq.heappush(unexplored_nodes, (g_scores[neighbor] + manhattan_heuristic(neighbor, target), neighbor))


    path = [target]
    while path[0] != start:
        path.insert(0, prev[path[0]])

    # print(" -> ".join([str(node) for node in path]))
    return g_scores[target]

part_1 = a_star(
    start=(0, 0),
    target=(GRID_SHAPE[0] - 1, GRID_SHAPE[1] - 1),
    grid=grid
)
print(f"Part 1: {part_1:,d}")

# Generate bigger grid
full_grid = []
for rep_col_idx in range(5):
    rep_row = []
    for rep_row_idx in range(5):
        new_subgrid = grid + rep_row_idx + rep_col_idx
        new_subgrid = np.where(new_subgrid > 9, new_subgrid + 1, new_subgrid) % 10
        rep_row.append(new_subgrid)
    full_grid.append(np.hstack(rep_row))
full_grid = np.vstack(full_grid)

part_2 = a_star(
    start=(0, 0),
    target=(full_grid.shape[0] - 1, full_grid.shape[1] - 1),
    grid=full_grid
)
print(f"Part 2: {part_2:,d}")