Advent of code 2024 day 7 jit 340 ms

@numba.njit
def day7_can_eval(desired, current, lst, part2):

  def count_digits(number):
    return 1 if number == 0 else math.floor(math.log10(abs(number))) + 1

  if current > desired:
    return False
  if not lst:
    return current == desired
  part2_candidate = current * 10 ** count_digits(lst[0]) + lst[0]
  return (
      day7_can_eval(desired, current * lst[0], lst[1:], part2)
      or day7_can_eval(desired, current + lst[0], lst[1:], part2)
      or (part2 and day7_can_eval(desired, part2_candidate, lst[1:], part2))
  )


def day7(s, part2=False):
  total = 0
  for line in s.splitlines():
    desired, *values = map(int, re.split(':? ', line))
    if day7_can_eval(desired, values[0], values[1:], part2):
      total += desired
  return total