advent day4B

from functools import reduce

with open("input", "r", encoding="utf-8") as f:
    data = [line.strip() for line in f.readlines()]

drawn_numbers = [int(nb) for nb in data[0].split(',')]


def parse_boards(data):
    boards = []
    count = -1
    for line in data[1:]:
        if not line:
            boards.append([])
            count += 1
        else:
            boards[count].append([[int(nb), False] for nb in line.split()])
    return boards


def mark_board(nb: int, board: list[list[list[int, bool]]]):
    for row in board:
        if [nb, False] in row:
            index = row.index([nb, False])
            row[index] = [nb, True]
            return
    return


def check_win(board: list[list[list[int, bool]]]) -> bool:
    len_row = len(board[0])
    for row in board:
        if sum(list(zip(*row))[1]) == len_row:
            return True
    transposed = zip(*board)
    for row in transposed:
        if sum(list(zip(*row))[1]) == len_row:
            return True
    return False


def play_game(boards: list[list[list[list[int, bool]]]], drawn: list[int]) -> int:
    has_won = [False] * len(boards)
    for nb in drawn:
        for i, board in enumerate(boards):
            if not has_won[i]:
                mark_board(nb, board)
                if check_win(board):
                    has_won[i] = True
                    if sum(has_won) == len(boards):
                        return score(board, nb)


def score(board: list[list[list[int, bool]]], nb: int) -> int:
    def add_unmarked(x, y): return x + (y[0] if not y[1] else 0)
    score = 0
    for row in board:
        score += reduce(add_unmarked, row, 0)
    return score * nb


boards = parse_boards(data)
print(play_game(boards, drawn_numbers))