Advent of code 2024 day 22 numba-jit 20 ms

@numba.njit  # Using numba with arrays allocated outside the loop.
def day22_jit(values, part2, base=19):
  part1_total = 0
  scores = np.full(base**4, 0)
  seen = np.full(base**4, False)
  secrets = np.full(2001, 0)

  for value in values:
    for i in range(2001):
      secrets[i] = value
      value = ((value << 6) ^ value) & 0xFFFFFF
      value = ((value >> 5) ^ value) & 0xFFFFFF
      value = ((value << 11) ^ value) & 0xFFFFFF

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

    seen[:] = False
    a = secrets[1] % 10 - secrets[0] % 10
    b = secrets[2] % 10 - secrets[1] % 10
    c = secrets[3] % 10 - secrets[2] % 10
    for i in range(2000 - 3):
      d = secrets[i + 4] % 10 - secrets[i + 3] % 10
      index = a * base**3 + b * base**2 + c * base**1 + d
      if not seen[index]:
        seen[index] = True
        scores[index] += secrets[i + 4] % 10
      a, b, c = b, c, d

  return scores.max() if part2 else part1_total