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) ] ) # (round, score) for each board (-1 when not found) winning_rounds: list[tuple[int, int]] = [(-1, -1) for _ in range(len(boards))] marked = np.zeros_like(boards, dtype=bool) for round, number in enumerate(numbers): # mark boards marked[boards == number] = True # check each board for winning 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]])), ) # all boards are winning - break if np.all(marked.all(axis=1) | marked.all(axis=2)): break # part 1 (_, 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}")