Advent of code 2024 day 6 faster

def day6(s, part2=False):
  grid = np.array([list(line) for line in s.splitlines()])
  ((y0, x0),) = np.argwhere(grid == '^')
  big = max(grid.shape)
  dydx_from_dir = [(0, -1), (-1, 0), (0, 1), (1, 0)]

  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.
    else:
      y, x = y1, x1

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

  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

  grid[y0, x0] = '^'
  count = 0
  for (obstacle_y, obstacle_x), ch in np.ndenumerate(grid):
    if ch == 'X':  # Candidate obstacle locations must lie on the original path.
      y, x, dir = y0, x0, 1
      visited = set()
      while True:
        dy, dx = dydx_from_dir[dir]
        steps = min(
          (jump_steps[dir, y, x],)
          + ((abs(d) - 1,) if obstacle_y == y and (d := obstacle_x - x) * dx > 0 else ())
          + ((abs(d) - 1,) if obstacle_x == x and (d := obstacle_y - y) * dy > 0 else ())
        )
        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  # Rotate clockwise.

  return count