Advent of code 2024 day 16

def day16(s, *, part2=False):
  grid = np.array([list(line) for line in s.splitlines()])
  ((y0, x0),) = np.argwhere(grid == 'S')
  ((yl, xl),) = np.argwhere(grid == 'E')
  state = y0, x0, 0, 1  # y, x, dy, dx.
  distance = collections.defaultdict(lambda: 10**9)  # For Dijkstra.
  distance[state] = 0
  pq = [(0, state)]

  def neighbors(d, state, sign=1):
    y, x, dy, dx = state
    yield d + sign, (y + dy * sign, x + dx * sign, dy, dx)
    yield d + 1000 * sign, (y, x, -dx, dy)
    yield d + 1000 * sign, (y, x, dx, -dy)

  while pq:
    d, state = heapq.heappop(pq)

    if state[:2] == (yl, xl):

      def visit_paths(state):
        d = distance.pop(state)  # Deletion for speedup.
        grid[state[:2]] = 'O'
        for d, state in neighbors(d, state, sign=-1):
          if d == distance[state]:
            visit_paths(state)

      visit_paths(state)
      return (grid == 'O').sum() if part2 else d

    for d, state in neighbors(d, state):
      if grid[state[:2]] != '#' and d < distance[state]:
        distance[state] = d
        heapq.heappush(pq, (d, state))