aoc-2021/4.py

from util import get_input, transpose
from more_itertools import grouper, flatten

input = get_input("4.input")

numbers = [int(a) for a in input[0].split(",")]

boards = [[[int(a) for a in row.split()] for row in c[1:]] for c in grouper(input[1:], 6, [])]

def check_board(numbers, board):
    for row in board + transpose(board):
        if len([a for a in row if a in numbers]) == 5:
            return True
    return False

def winners(numbers, boards):
    all_winners = []
    for i in range(len(numbers)):
        winners = [b for b in boards if check_board(numbers[:1+i], b)]
        for board in winners:
            boards.remove(board)
            unmarked = [a for a in flatten(board) if a not in numbers[:1+i]]
            all_winners.append(sum(unmarked) * numbers[i])
    return all_winners

all_winners = winners(numbers, boards)
print(all_winners[0])
print(all_winners[-1])