advent_of_code_2021/day04/part2.py

import argparse
from typing import List, Set
import pprint
import sys

parser = argparse.ArgumentParser()
parser.add_argument("ifile", type=argparse.FileType('r'))
args = parser.parse_args()
lines = [line.replace("  ", " ").strip() for line in args.ifile.readlines()]

number_calls = [num for num in lines[0].strip().split(',')]

boards = [[line.split(' ') for line in lines[x:x+5]] for x in range(2, len(lines), 6)]


def winning_board(board: List[List[str]], called_numbers: Set[str]) -> bool:
    for i in range(0, 5):
        if all([board[i][col] in called_numbers for col in range(0, 5)]):
            return True
        if all([board[row][i] in called_numbers for row in range(0, 5)]):
            return True
    return False


def score_board(board: List[List[str]], called_numbers: Set[str]) -> int:
    return sum([int(num) for row in board for num in row if num not in called_numbers])


called_numbers: Set[str] = set()
most_recent = None

boards_that_have_won: Set[int] = set()

for number in number_calls:
    called_numbers.add(number)

    for index, board in enumerate(boards):
        if not (index in boards_that_have_won) and winning_board(board, called_numbers):
            boards_that_have_won.add(index)
            most_recent = score_board(board, called_numbers) * int(number)

print(most_recent)