Advent of code 2024 day 10 numba fastest

@numba.njit
def day10_jit(grid, part2):
  neighbors = (0, -1), (0, 1), (-1, 0), (1, 0)
  total = 0

  if not part2:
    for y0 in range(1, grid.shape[0] - 1):
      for x0 in range(1, grid.shape[1] - 1):
        if grid[y0, x0] == 0:
          active = {(y0, x0)}  # Set of positions reached after `dist` moves.
          for dist in range(1, 10):
            active2 = set()
            for y, x in active:
              for dy, dx in neighbors:
                y2, x2 = y + dy, x + dx
                if grid[y2, x2] == dist:
                  active2.add((y2, x2))
            active = active2
          total += len(active)

  else:
    a = np.empty_like(grid)
    for y in range(1, grid.shape[0] - 1):
      for x in range(1, grid.shape[1] - 1):
        if grid[y, x] == 9:
          a[y, x] = 1
    for dist in range(8, -1, -1):
      for y in range(1, grid.shape[0] - 1):
        for x in range(1, grid.shape[1] - 1):
          if grid[y, x] == dist:
            count = 0
            for dy, dx in neighbors:
              y2, x2 = y + dy, x + dx
              if grid[y2, x2] == dist + 1:
                count += a[y2, x2]
            a[y, x] = count
    for y in range(1, grid.shape[0] - 1):
      for x in range(1, grid.shape[1] - 1):
        if grid[y, x] == 0:
          total += a[y, x]

  return total


def day10(s, part2=False):
  grid = np.array([list(line) for line in s.splitlines()], int)
  grid = np.pad(grid, 1, constant_values=-1)
  return day10_jit(grid, part2)