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}")