Advent of code 2024 day 22 numpy fully vectorized

def day22c(s, *, part2=False, window_size=4, base=19):
  vec = np.array(s.splitlines(), int)
  n_inputs = len(vec)

  secret = np.empty((n_inputs, 2001), int)
  for i in range(2001):
    secret[:, i] = vec
    vec = ((vec << 6) ^ vec) & 0xFFFFFF
    vec = ((vec >> 5) ^ vec) & 0xFFFFFF
    vec = ((vec << 11) ^ vec) & 0xFFFFFF

  if not part2:
    return secret[:, -1].sum()

  digit = secret % 10
  diffs2 = np.diff(digit, axis=1) + 9
  score = digit[:, 4:]

  # Create sliding windows along axis 1; shape=[n_inputs, 1997, 4].
  windows = np.lib.stride_tricks.sliding_window_view(diffs2, window_shape=window_size, axis=1)

  # Flatten each window into a scalar.
  windows_flat = np.einsum('ijk,k->ij', windows, base ** np.arange(4))

  # Create a mask identifying all the first occurrences in each row.
  index = np.argsort(windows_flat, axis=1, kind='stable')
  sorted_windows = np.take_along_axis(windows_flat, index, axis=1)
  is_different = np.diff(sorted_windows, axis=1, prepend=-1) != 0
  mask = np.empty_like(windows_flat, bool)
  np.put_along_axis(mask, index, is_different, axis=1)

  # Accumulate values into the totals array.
  rows, cols = mask.nonzero()
  flat_indices = windows_flat[rows, cols]
  values = score[rows, cols]

  totals = np.zeros(base**4, int)
  np.add.at(totals, flat_indices, values)
  return totals.max()