Advent of code 2024 day 22 numba-jit 370 ms

@numba.njit
def day22_jit(values, part2):
  part1_total = 0
  scores = {}  # Unfortunately, numba does not support defaultdict(int).
  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 = set()
    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
      seq = a, b, c, d
      if seq not in seen:
        seen.add(seq)
        if seq in scores:
          scores[seq] += secrets[i + 4] % 10
        else:
          scores[seq] = secrets[i + 4] % 10
      a, b, c = b, c, d

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


def day22(s, *, part2=False):
  values = list(map(int, s.splitlines()))
  return day22_jit(values, part2)