from itertools import product, chain from utils import open_day class Board: cells: list[int | None] width: int height: int col_hits: list[int] row_hits: list[int] has_bingo: bool unmarked_sum: int def __init__(self, cells: list[list[int]]): self.width = len(cells[0]) self.height = len(cells) self.col_hits = [0] * self.width self.row_hits = [0] * self.height self.unmarked_sum = 0 self.cells = [] for cell in chain.from_iterable(cells): self.unmarked_sum += cell self.cells.append(cell) self.has_bingo = False def call(self, num: int): for y, x in product(range(self.height), range(self.width)): cell = self.cells[x + self.width * y] if cell == num: self.col_hits[x] += 1 self.row_hits[y] += 1 self.unmarked_sum -= cell self.has_bingo = ( any(hits == self.height for hits in self.col_hits) or any(hits == self.width for hits in self.row_hits) ) return @staticmethod def from_string(s: str) -> 'Board': return Board([[int(n) for n in l.split()] for l in s.split('\n')]) def solve(nums: list[int], boards: list[Board]) -> tuple[int, int]: won: set[int] = set() wins: list[int] = [] num: int for num in nums: i: int board: Board for i, board in enumerate(boards): if i in won: continue board.call(num) if board.has_bingo: won.add(i) wins.append(num * board.unmarked_sum) return wins[0], wins[-1] nums: str | list[int] boards: list[str] | list[Board] nums, *boards = open_day(4).read().rstrip().split('\n\n') nums = list(map(int, nums.split(','))) boards = [Board.from_string(b) for b in boards] part1, part2 = solve(nums, boards) print(part1) print(part2)