Advent of code 2024 day 6 jit fast 12 ms

@numba.njit
def day6_jit(grid: np.ndarray, y0: int, x0: int, part2: bool):
  y, x, dy, dx = y0, x0, -1, 0
  count = 0
  while True:
    if grid[y, x] != 'X':
      count += 1
    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 count

  big = max(grid.shape)
  jump_steps = np.empty((4, *grid.shape), np.int32)
  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] = '^'
  dydx_from_dir = [(0, -1), (-1, 0), (0, 1), (1, 0)]
  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 = jump_steps[dir, y, x]
        if obstacle_y == y:
          d = obstacle_x - x
          if d * dx > 0:
            steps = min(steps, abs(d) - 1)
        if obstacle_x == x:
          d = obstacle_y - y
          if d * dy > 0:
            steps = min(steps, abs(d) - 1)
        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


def day6(s, part2=False):
  grid = np.array([list(line) for line in s.splitlines()])
  ((y0, x0),) = np.argwhere(grid == '^')
  return day6_jit(grid, y0, x0, part2)