Advent of code 2024 day 10

def day10(s, part2=False):
  grid = np.array([list(line) for line in s.splitlines()], int)
  grid = np.pad(grid, 1, constant_values=-1)
  trailheads = np.argwhere(grid == 0)
  neighbors = ((0, -1), (0, 1), (-1, 0), (1, 0))
  total = 0

  for y0, x0 in trailheads:
    if not part2:
      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:
      active = collections.Counter([(y0, x0)])  # Number of paths reaching position.
      for dist in range(1, 10):
        active2 = collections.Counter()
        for (y, x), count in active.items():
          for dy, dx in neighbors:
            y2, x2 = y + dy, x + dx
            if grid[y2, x2] == dist:
              active2[y2, x2] += count
        active = active2
      total += active.total()

  return total