Advent of code 2023 day 6

def day6(s, *, part2=False):
  if part2:
    s = s.replace(' ', '')
  ts, ds = (map(int, line.split(':')[1].split()) for line in s.splitlines())

  def num_wins(t, d):
    # Quadratic equation -t**2 * x**2 + t * x - d = 0 has roots x = u \pm v with:
    u, v = t / 2, (t**2 - 4 * d) ** 0.5 / 2
    return math.ceil(u + v - 1) - math.floor(u - v + 1) + 1

  return math.prod(num_wins(t, d) for t, d in zip(ts, ds))