Advent of code 2023 day 17

def day17(s, *, part2=False, visualize=False, rep=2):  # Dijkstra search over "larger" state node.
  grid = np.array([list(line) for line in s.splitlines()], int)

  state = 0, 0, 0, 1, 0  # (y, x, dy, dx, dn).
  target_yx = grid.shape[0] - 1, grid.shape[1] - 1
  distances = collections.defaultdict(lambda: 10**9)
  distances[state] = 0
  priority_queue = [(0, state)]
  prev = {}

  while priority_queue:
    _, state = heapq.heappop(priority_queue)
    y, x, dy, dx, dn = state
    distance = distances[state]
    if (y, x) == target_yx and (not part2 or dn >= 4):
      break
    for dy2, dx2, same_dir in [(-dx, dy, False), (dy, dx, True), (dx, -dy, False)]:
      dn2 = dn + 1 if same_dir else 1
      if (dn2 <= 10 and (dn == 0 or dn >= 4 or same_dir)) if part2 else (dn2 <= 3):
        y2, x2 = y + dy2, x + dx2
        if 0 <= y2 < grid.shape[0] and 0 <= x2 < grid.shape[1]:
          state2 = y2, x2, dy2, dx2, dn2
          distance2 = distance + grid[y2, x2]
          if distance2 < distances[state2]:
            distances[state2] = distance2
            prev[state2] = state
            heapq.heappush(priority_queue, (distance2, state2))
  else:
    raise ValueError('No solution found.')

  if visualize:
    path = [state[:2]]
    while state[:2] != (0, 0):
      state = prev[state]
      path.append(state[:2])
    assert sum(grid[yx] for yx in path[:-1]) == distance

    image = media.to_rgb(grid * 0.9, cmap='gray')  # Or cmap='bwr'.
    image2 = image.copy()
    image2[tuple(np.array(path).T)] = 255, 255, 0
    image = image.repeat(rep, 0).repeat(rep, 1)
    image2 = image2.repeat(rep, 0).repeat(rep, 1)
    title = 'day17b' if part2 else 'day17a'
    media.show_video([image, image2], codec='gif', fps=1, title=title)

  return distance