Advent of code 2024 day 6 jit fast 7 ms

# For Part 2, (1) precompute a jump map and (2) march along the path and look for
# a possible loop only when encountering an empty space for the first time.
@numba.njit
def day6_jit(grid: np.ndarray, y0: int, x0: int, part2: bool) -> int:
  big = max(grid.shape)
  dydx_from_dir = [(0, -1), (-1, 0), (0, 1), (1, 0)]
  jump_steps = np.empty((4, *grid.shape), np.int32)
  if part2:
    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

  def adding_obstacle_would_create_loop(y, x, dir, obstacle_y, obstacle_x):
    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:
        return False
      y, x = y + dy * steps, x + dx * steps
      state = y, x, dir
      if state in visited:
        return True
      visited.add(state)
      dir = (dir + 1) % 4  # Rotate clockwise.

  y, x, dir = y0, x0, 1
  count = 0
  while True:
    if not part2 and grid[y, x] != 'X':
      grid[y, x] = 'X'
      count += 1
    dy, dx = dydx_from_dir[dir]
    y1, x1 = y + dy, x + dx
    if not (0 <= y1 < grid.shape[0] and 0 <= x1 < grid.shape[1]):
      break
    if grid[y1, x1] == '#':
      dir = (dir + 1) % 4  # Rotate clockwise.
    else:
      if part2 and grid[y1, x1] == '.':
        grid[y1, x1] = 'Y'
        count += adding_obstacle_would_create_loop(y, x, dir, y1, x1)
      y, x = y1, x1

  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)