Advent of code 2023 day 11 fast

def day11(s, *, part2=False):
  growth = 999_999 if part2 else 1
  grid = np.array([list(line) for line in s.splitlines()])
  galaxies = np.argwhere(grid == '#')
  all_g1_g2 = itertools.combinations(galaxies, 2)
  total = sum(abs(a - b) for g1, g2 in all_g1_g2 for a, b in zip(g1, g2))
  for c in range(2):
    for i in range(len(grid)):
      if np.all(grid[i] == '.'):
        num_before = sum(1 for galaxy in galaxies if galaxy[c] < i)
        total += growth * num_before * (len(galaxies) - num_before)
    grid = grid.T
  return total