advent_15_part2.py

import heapq

def get(x, y, cave):
  overflow = y // len(cave)
  y = y % len(cave)
  overflow += x // len(cave[y])
  x = x % len(cave[y])
  value = cave[y][x]
  value += overflow
  value %= 9
  if value == 0:
    value = 9
  return value

def adjacent(x, y, cave):
  max_y = len(cave) * 5
  max_x = len(cave[0]) * 5
  deltas = [(-1, 0), (0, -1), (1, 0), (0, 1)]
  for delta_x, delta_y in deltas:
    candidate_x = x + delta_x
    candidate_y = y + delta_y
    if (candidate_y >= 0 and candidate_y < max_y and
        candidate_x >= 0 and candidate_x < max_x):
      yield candidate_x, candidate_y
  return

def key(x, y):
  return "{0}-{1}".format(x, y)

def dijkstra(position_x, position_y, target_x, target_y, cave):
  visited = set()
  heap = [(0, position_x, position_y)]
  while heap:
    cost, x, y = heapq.heappop(heap)
    if x == target_x and y == target_y:
      return cost
    current_key = key(x, y)
    if current_key in visited:
      continue
    visited.add(current_key)
    for new_x, new_y in adjacent(x, y, cave):
      heapq.heappush(heap, (cost + get(new_x, new_y, cave), new_x, new_y))
  return - 1

if __name__ == '__main__':
  with open('input.txt') as data:
    cave = []
    for line in data:
      line = line.strip()
      cave.append([int(risk) for risk in line])
  print(dijkstra(0, 0, len(cave[0]) * 5 - 1, len(cave) * 5 - 1, cave))