Advent of code 2024 day 6

def day6(s, part2=False):  # Faster with precomputed jumps.
  grid = np.array([list(line) for line in s.splitlines()])
  ((y0, x0),) = np.argwhere(grid == '^')
  big = max(grid.shape)

  y, x, dy, dx = y0, x0, -1, 0
  while True:
    grid[y, x] = 'X'
    y1, x1 = y + dy, x + dx
    if not (0 <= y1 < grid.shape[0] and 0 <= x1 < grid.shape[1]):
      break
    if grid[y1, x1] == '#':
      dy, dx = dx, -dy  # Rotate clockwise.
      continue
    y, x = y1, x1

  if not part2:
    return (grid == 'X').sum()

  grid[y0, x0] = '^'
  dydx_from_dir = [(0, -1), (-1, 0), (0, 1), (1, 0)]

  jump_steps = np.empty((4, *grid.shape), int)
  for dir in range(4):
    rotated_grid = np.rot90(grid, k=dir)
    rotated_jump_steps = np.rot90(jump_steps[dir], k=dir)
    for y, row in enumerate(rotated_grid):
      num = big
      for x, ch in enumerate(row):
        num = -1 if ch == '#' else num + 1
        rotated_jump_steps[y, x] = num

  count = 0
  for obstacle_yx, ch in np.ndenumerate(grid):
    if ch == 'X':  # Candidate obstacle locations must lie on the original path.
      y, x, dir = y0, x0, 1
      grid[obstacle_yx] = '#'
      visited = set()
      while True:
        dy, dx = dydx_from_dir[dir]
        if obstacle_yx[0] == y or obstacle_yx[1] == x:
          y1, x1 = y + dy, x + dx
          if not (0 <= y1 < grid.shape[0] and 0 <= x1 < grid.shape[1]):
            break
          if grid[y1, x1] != '#':
            y, x = y1, x1
            continue
        else:
          steps = jump_steps[dir, y, x]
          if steps >= big:
            break
          y, x = y + dy * steps, x + dx * steps
        state = y, x, dir
        if state in visited:
          count += 1
          break
        visited.add(state)
        dir = (dir + 1) % 4

      grid[obstacle_yx] = 'X'  # Remove the temporarily added obstacle.

  return count