Advent of code 2024 day 22 all-generators

def day22(s, *, part2=False):

  def generate_secrets(secret):
    for _ in range(2001):
      yield secret
      secret = ((secret << 6) ^ secret) & 0xFFFFFF
      secret = ((secret >> 5) ^ secret) & 0xFFFFFF
      secret = ((secret << 11) ^ secret) & 0xFFFFFF

  def changes(secrets):
    return (b % 10 - a % 10 for a, b in itertools.pairwise(secrets))

  initials = map(int, s.splitlines())

  if not part2:
    return sum(more_itertools.last(generate_secrets(initial)) for initial in initials)

  totals: Any = collections.defaultdict(int)

  for initial in initials:
    seen = set()
    secrets = generate_secrets(initial)
    it1, it2 = itertools.tee(secrets)
    windows = more_itertools.windowed(changes(it1), 4)
    scores = (secret % 10 for secret in itertools.islice(it2, 4, None))
    for window, score in zip(windows, scores):
      if window not in seen:
        seen.add(window)
        totals[window] += score

  return max(totals.values())