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])