advent-of-code/2021/day4.py

41 lines
1.0 KiB
Python
Raw Normal View History

2023-12-09 10:01:24 +00:00
import sys
import numpy as np
lines = sys.stdin.read().splitlines()
numbers = [int(c) for c in lines[0].split(",")]
boards = np.asarray(
[
[[int(c) for c in line.split()] for line in lines[start : start + 5]]
for start in range(2, len(lines), 6)
]
)
# part 1
marked = np.zeros_like(boards, dtype=bool)
winning_rounds: list[tuple[int, int]] = [(-1, -1) for _ in range(len(boards))]
for round, number in enumerate(numbers):
marked[boards == number] = True
for index in range(len(boards)):
if winning_rounds[index][0] > 0:
continue
if np.any(np.all(marked[index], axis=0) | np.all(marked[index], axis=1)):
winning_rounds[index] = (
round,
number * int(np.sum(boards[index][~marked[index]])),
)
if np.all(marked.all(axis=1) | marked.all(axis=2)):
break
(_, score) = min(winning_rounds, key=lambda w: w[0])
print(f"answer 1 is {score}")
# part 2
(_, score) = max(winning_rounds, key=lambda w: w[0])
print(f"answer 2 is {score}")