#!/usr/bin/env python3 import sys from typing import Optional, Sequence class Board: def __init__(self, matrix: Sequence[Sequence[int]]): self.matrix = matrix self.marks = [[False] * len(matrix[0]) for _ in range(len(matrix[0]))] def mark(self, num: int): for i, row in enumerate(self.matrix): for j, col in enumerate(row): if col == num: self.marks[i][j] = True def row(self, i: int) -> Sequence[int]: return self.matrix[i] def col(self, j: int) -> Sequence[int]: return [self.matrix[i][j] for i in range(len(self.matrix))] def row_marks(self, i: int) -> Sequence[int]: return self.marks[i] def col_marks(self, j: int) -> Sequence[int]: return [self.marks[i][j] for i in range(len(self.marks))] def row_marked(self) -> Optional[int]: "If an entire row is marked, return that row's index. Else return None." for i in range(len(self.matrix)): if all(self.row_marks(i)): return i return None def col_marked(self) -> Optional[int]: "If an entire column is marked, return that column's index. Else return None." for j in range(len(self.matrix[0])): if all(self.col_marks(j)): return j return None def all_unmarked(self) -> Sequence[int]: "Gets a list of all the unmarked numbers on this board" return [ self.matrix[i][j] for i in range(len(self.matrix)) for j in range(len(self.matrix[i])) if not self.marks[i][j] ] def part1(draws: Sequence[int], matrices: Sequence[Sequence[Sequence[int]]]): boards = [Board(matrix) for matrix in matrices] result = None for draw in draws: # Mark the boards for board in boards: board.mark(draw) # Check to see if there are 5 in a row in any direction row_num = board.row_marked() if row_num is not None: result = sum(board.all_unmarked()) * draw break col_num = board.col_marked() if col_num is not None: result = sum(board.all_unmarked()) * draw break if result is not None: break print(result) def part2(draws: Sequence[int], matrices: Sequence[Sequence[Sequence[int]]]): boards = [Board(matrix) for matrix in matrices] result = None for draw in draws: # Mark the boards for board in boards: board.mark(draw) # Filter out the winners if len(boards) == 1: board = boards[0] if board.row_marked() is not None or board.col_marked() is not None: result = sum(board.all_unmarked()) * draw break else: boards = [ board for board in boards if board.row_marked() is None and board.col_marked() is None ] print(result) lines = [line.strip() for line in sys.stdin if line.strip()] draws = [int(draw) for draw in lines.pop(0).split(",")] # Matrices are 5 lines matrices = [ [[int(col) for col in line.split()] for line in lines[i : i + 5]] for i in range(0, len(lines), 5) ] print("Part 1") part1(draws, matrices) print("Part 2") part2(draws, matrices)