aoc2024day22.py

# === Parsing ===

with open("input.txt", "r") as f:
    seeds = [int(line) for line in f]


# === Part 1 ===
from collections import deque
from itertools import islice


def monkey_prng(seed):
    value = seed
    mod = pow(2, 24)
    while True:
        value = ((value << 6) ^ value) % mod
        value = ((value >> 5) ^ value) % mod
        value = ((value << 11) ^ value) % mod
        yield value


total_price = sum(
    deque(islice(monkey_prng(seed), 2000), maxlen=1).pop() for seed in seeds
)
print(total_price)  # 16039090236


# === Part 2 ===
from itertools import chain, pairwise, tee

prices = [
    [p % 10 for p in chain([seed] + list(islice(monkey_prng(seed), 2000)))]
    for seed in seeds
]
deltas = [[b - a for a, b in pairwise(price_list)] for price_list in prices]


def fourwise(sequence):
    a, b, c, d = tee(sequence, 4)
    next(b)
    next(c)
    next(c)
    next(d)
    next(d)
    next(d)
    yield from zip(a, b, c, d)


seqs_to_prices_list = []
all_seqs = set()
# For each monkey, we create a dictionary that maps sequences to their
# prices in bananas.
for price_list, delta_list in zip(prices, deltas):
    seq_to_price_mapping = {}
    for idx, seq in enumerate(fourwise(delta_list)):
        if seq not in seq_to_price_mapping:
            seq_to_price_mapping[seq] = price_list[idx + 4]
            all_seqs.add(seq)
    seqs_to_prices_list.append(seq_to_price_mapping)

print(len(all_seqs))
best_price = 0
best_seq = 0
for seq in all_seqs:
    bananas = 0
    for seq_to_price_mapping in seqs_to_prices_list:
        bananas += seq_to_price_mapping.get(seq, 0)
    if bananas > best_price:
        best_price = bananas
        best_seq = seq

print(best_seq, best_price)