advent_of_code_2021/day04/part1.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()
for number in number_calls:
    called_numbers.add(number)

    for board in boards:
        if winning_board(board, called_numbers):
            print(score_board(board, called_numbers) * int(number))
            sys.exit(0)