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

122 lines
3 KiB
Python

#!/usr/bin/env python3
import sys
import numpy as np
import rich.progress
input_file = sys.argv[1]
with open(input_file) as fd:
lines = [line.rstrip() for line in fd.readlines()]
buyers = len(lines)
def prng(secwet: int) -> int:
secwet ^= secwet * 64
secwet %= 16777216
secwet ^= secwet // 32
secwet %= 16777216
secwet ^= secwet * 2048
secwet %= 16777216
return secwet
buyers_ban: list[list[int]] = []
buyers_dif: list[list[int]] = []
for line in rich.progress.track(lines, description="Calculating future banana prices"):
secwet = int(line)
buyer_ban: list[int] = []
buyer_dif: list[int] = []
last_ban = secwet % 10
for i in range(2000):
secwet = prng(secwet)
ban = secwet % 10
buyer_ban.append(ban)
dif = ban - last_ban
buyer_dif.append(dif)
last_ban = ban
# print(f"{secwet=} {ban=} {dif=}")
# print(f"{buyer_ban=}")
# print(f"{buyer_dif=}")
buyers_ban.append(buyer_ban)
buyers_dif.append(buyer_dif)
buyers_dif_np = np.array(buyers_dif)
buyers_ban_np = np.array(buyers_ban)
sequence = tuple[int, int, int, int]
def totbans(seq: sequence) -> int:
match = None
for i, num in enumerate(seq):
nmatch = buyers_dif_np == num
if match is not None:
# nmatch = sp.ndimage.shift(nmatch, (0, -i))
nmatch = np.roll(nmatch, -i, axis=1)
nmatch &= match
match = nmatch
# bans = buyers_ban_np * match
# found = match.max(axis=1)
# indexes = np.argmax(match > 0, axis=1)
tot = 0
assert match is not None
for b, buyer_match in enumerate(match):
if not buyer_match.max():
continue
arg: int = np.argmax(buyer_match > 0)
try:
ban = buyers_ban_np[b, arg+3]
tot += ban
except IndexError:
pass # shrug
return tot
for buyer in range(buyers):
buyer_dif = buyers_dif[buyer]
for i in range(2000 - 4):
if (
buyer_dif[i] == seq[0]
and buyer_dif[i + 1] == seq[1]
and buyer_dif[i + 2] == seq[2]
and buyer_dif[i + 3] == seq[3]
):
# if tuple(buyer_dif[i : i + 4]) == seq:
tot += buyers_ban[buyer][i + 3]
break
return tot
# print(f"{totbans((6, -1, -1, 0))=}") # demo0
# print(f"{totbans((-2, 1, -1, 3))=}") # demo2 aoc
# print(f"{totbans((6, -4, 4, -9))=}") # demo2 first
all_seqs: set[sequence] = set()
for buyer in rich.progress.track(
range(buyers), description="Generating possible sequences"
):
buyer_dif = buyers_dif[buyer]
for i in range(2000 - 4):
seq: sequence = tuple(buyer_dif[i : i + 4])
all_seqs.add(seq)
print(f"{len(all_seqs)=}")
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)
# 1909 too low