advent2021/four_second.py

29 lines
873 B
Python
Executable File

#!/usr/bin/env python3
from useful import *
def win(b, d):
hit = np.vectorize(d.get)(b)
size = b.shape[0]
return any(correct == size for correct in (*hit.sum(axis=0), *hit.sum(axis=1)))
def score(b, d, last):
return last * sum(x for x in b.reshape(np.prod(b.shape)) if not d[x])
with open(0) as f:
header, footer = f.read().split('\n\n', maxsplit=1)
pool = [int(x) for x in header.strip().split(',')]
boards = [np.asarray([int(n) for n in board.strip().split()]).reshape((5,5))
for board in footer.split('\n\n')]
drawn = {n: False for n in pool} # to vectorize
winners = []
for n in pool:
drawn[n] = True
for i in range(len(boards))[::-1]:
b = boards[i]
if win(b, drawn):
winners.append(score(b, drawn, n))
del boards[i]
print(winners[-1])