gregs-aoc-2021-day-04

#!/usr/bin/env python3

SAMPLE_MODE = False
PART_2 = True

class BingoBoard:
    def __init__(self, nums):
        self.nums = nums
        assert len(self.nums) == 25
        self.marked = [False] * 25
        self.won = False

    def mark_num_if_available(self, num):
        if num in self.nums:
            index = self.nums.index(num)
            assert not self.marked[index]
            self.marked[index] = True

    def mark_won_if_winner(self):
        for i in range(5):
            if self._row_is_a_winner(i) or self._col_is_a_winner(i):
                self.won = True
                return
    
    def _row_is_a_winner(self, row_num):
        return all([self.marked[row_num*5 + i] for i in range(5)])
    
    def _col_is_a_winner(self, col_num):
        return all([self.marked[i*5 + col_num] for i in range(5)])

    def score(self, most_recently_called_number):
        return self._sum_of_unmarked_squares() * most_recently_called_number

    def _sum_of_unmarked_squares(self):
        return sum([int(self.nums[i]) for i in range(25) if not self.marked[i]])


class NoWinnerError(Exception):
    pass


def load_input(sample=False):
    filename = "sample_input.txt" if sample else "input.txt"
    with open(filename) as f:
        return [line.strip() for line in f.readlines()]

def parse_input(lines):
    called = [int(n) for n in lines[0].split(",")]
    boards = []
    for i in range(int((len(lines) - 1) / 6)):
        board_lines = lines[i*6 + 2 : i*6 + 7]
        board_nums = [int(n) for n in " ".join(board_lines).split()]
        boards.append(BingoBoard(board_nums))
    return called, boards

def play_bingo(nums, boards):
    one_board_remains = False
    for num in nums:
        print()
        print(f"Calling {num}")
        if PART_2 and len([board for board in boards if not board.won]) == 1:
            print("One board remains!")
            one_board_remains = True
        for board in boards:
            if PART_2 and board.won:
                continue
            print(f"Parsing board {board.nums}")
            board.mark_num_if_available(num)
            board.mark_won_if_winner()
            if (not PART_2) and board.won:
                return board.score(num)
            elif PART_2 and one_board_remains and board.won:
                return board.score(num)
    raise NoWinnerError()

if __name__ == "__main__":
    lines = load_input(SAMPLE_MODE)
    nums, boards = parse_input(lines)
    print(play_bingo(nums, boards))