advent-of-code/2024/22/two_fast.py
2024-12-25 12:59:49 +01:00

113 lines
3.1 KiB
Python

#!/usr/bin/env python3
import collections
import math
import sys
import numpy as np
import rich.progress
dtype = np.int32
input_file = sys.argv[1]
with open(input_file) as fd:
buyers_init_secret = np.array([int(line) for line in fd.readlines()], dtype=dtype)
buyers = len(buyers_init_secret)
ITERATIONS = 2000
SEQUENCE_LENGTH = 4
buyers_prices = np.ndarray((buyers, ITERATIONS), dtype=dtype)
buyers_diffs = np.ndarray((buyers, ITERATIONS), dtype=dtype)
buyers_secret = buyers_init_secret.copy()
last_buyers_price = buyers_init_secret % 10
for i in rich.progress.track(range(ITERATIONS), description="Simulating stock market"):
buyers_secret ^= buyers_secret * 64
buyers_secret %= 16777216
buyers_secret ^= buyers_secret // 32
buyers_secret %= 16777216
buyers_secret ^= buyers_secret * 2048
buyers_secret %= 16777216
buyers_price = buyers_secret % 10
buyers_diff = buyers_price - last_buyers_price
buyers_prices[:, i] = buyers_price
buyers_diffs[:, i] = buyers_diff
last_buyers_price = buyers_price
# Compress sequence tuples into a single integer
# Move to positive numbers
trans = -buyers_diffs.min()
buyers_diffs_translated = buyers_diffs + trans
# Decide on a value to shift
shift = math.ceil(math.log2(buyers_diffs_translated.max()))
buyers_sequences = buyers_diffs_translated.copy()
for i in range(1, SEQUENCE_LENGTH):
buyers_sequences += np.roll(buyers_diffs_translated << (shift * i), i, axis=1)
# Make first few sequences invalid
buyers_sequences[:, :SEQUENCE_LENGTH] = -1
# NEW (knowing it's best to iterate per buyer than per sequence) ~ 5 seconds
# Inspired by the following, which is even faster, probably because data locality
# achieves more than SIMD-fying?
# https://github.com/mkern75/AdventOfCodePython/blob/23b6becdc873c6b865e783122a7dbce0b5f40f60/year2024/Day22.py
max_sequence = 1 << shift * SEQUENCE_LENGTH
sequences_total = np.zeros(max_sequence, dtype=dtype)
seen = np.zeros((buyers, max_sequence), dtype=bool)
for b in rich.progress.track(range(buyers), description="Evaluating each buyer"):
for i in range(SEQUENCE_LENGTH, ITERATIONS):
seq = buyers_sequences[b, i]
if seen[b, seq]:
continue
seen[b, seq] = True
sequences_total[seq] += buyers_prices[b, i]
print(f"{sequences_total.argmax()=}")
print(sequences_total.max())
sys.exit(0)
# OLD (knowing you can compresses sequences only) ~ 1.5 minute
def totbans(seq: int) -> int:
match = buyers_sequences == seq
found = match.max(axis=1)
indexes = np.argmax(match, axis=1)
bans = buyers_prices[range(buyers), indexes]
bans *= found
return bans.sum()
def seq_to_int(seq: tuple[int, ...]) -> int:
tot = 0
for s, num in enumerate(seq):
tot += (num + trans) << (SEQUENCE_LENGTH - s - 1) * shift
return tot
print(f"{totbans(seq_to_int((-2, 1, -1, 3)))=}")
all_seqs: set[int] = set(buyers_sequences.flat) - {-1}
maxi = 0
max_seq = None
for seq in rich.progress.track(all_seqs, description="Finding score for sequences"):
tb = totbans(seq)
if tb > maxi:
maxi = tb
max_seq = seq
print(f"{max_seq=}")
print(maxi)