AdventOfCode/2021/day4/day4_part2_v2.py

14 lines
282 B
Python

import numpy as np
import sys
n, *b = open(sys.argv[1])
b = np.loadtxt(b,int).reshape(-1,5,5)
for n in map(int, n.split(',')):
b[b == n] = -1
m = (b == -1)
win = (m.all(1) | m.all(2)).any(1)
if win.any():
print((b * ~m)[win].sum() * n)
b = b[~win]