Advent of code 2024 day 22 concise

def day22(s, *, part2=False):
  total = 0
  scores = collections.defaultdict(int)
  secrets = np.empty(2001, int)

  for v in map(int, s.splitlines()):
    for i in range(2001):
      secrets[i] = v
      v = ((v << 6) ^ v) & 0xFFFFFF
      v = ((v >> 5) ^ v) & 0xFFFFFF
      v = ((v << 11) ^ v) & 0xFFFFFF

    if not part2:
      total += secrets[-1]
      continue

    seen = set()
    diffs = np.diff(secrets % 10)
    for seq, bananas in zip(more_itertools.windowed(diffs, 4), secrets[4:]):
      if seq not in seen:
        seen.add(seq)
        scores[seq] += bananas % 10

  return max(scores.values()) if part2 else total