from itertools import chain from utils import open_day class Board: cells: list[int] width: int height: int col_hits: list[int] row_hits: list[int] has_bingo: bool unmarked_sum: int nums: set[int] def __init__(self, cells: list[list[int]]) -> None: 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 = [] self.nums = set() for cell in chain.from_iterable(cells): self.unmarked_sum += cell self.cells.append(cell) self.nums.add(cell) self.has_bingo = False def call(self, num: int) -> None: if num not in self.nums: return pos: int = self.cells.index(num) x: int = pos % self.width y: int = pos // self.width self.col_hits[x] += 1 self.row_hits[y] += 1 self.unmarked_sum -= num self.has_bingo = ( self.has_bingo or self.col_hits[x] == self.height or self.row_hits[y] == self.width ) @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)