adventofcode2021/aoc4-2.py

class Board:
    def __init__(self, lines, boardsize=None, boardnum=None):
        self.lines = [[el for el in line.split(" ") if el != ""] for line in lines]
        self.boardsize = 5 if boardsize == None else boardsize
        self.boardnum = 0 if boardnum == None else boardnum
        self.marks = []
        for _ in range(self.boardsize):
            self.marks.append([0]*self.boardsize)

    def mark(self, digit):
        import pdb
        # if self.boardnum == 1:
        #     pdb.set_trace()
        for i, line in enumerate(self.lines):
            if digit in line:
                y = i
                x = line.index(digit)
                self.marks[y][x] = 1
                return None

    def check(self):
        for line in self.marks:
            if sum(line) == 5:
                return True
        for i in range(5):
            if  self.marks[0][i] == 1 and \
                self.marks[1][i] == 1 and \
                self.marks[2][i] == 1 and \
                self.marks[3][i] == 1 and \
                self.marks[4][i] == 1:
                return True
        return False

    def sum_unmarked(self):
        total = 0
        for i in range(self.boardsize):
            for j in range(self.boardsize):
                if self.marks[j][i] == 0:
                    total += int(self.lines[j][i])
        return total

    def __repr__(self):
        outstr = f"Board {self.boardnum}:\n"
        for i in range(self.boardsize):
            outstr += f"{' '.join([str(el).rjust(2) for el in self.lines[i]])}   {' '.join([str(el).rjust(2) for el in self.marks[i]])}\n"
        return outstr

    def __str__(self):
        outstr = f"Board {self.boardnum}:\n"
        for i in range(self.boardsize):
            outstr += f"{' '.join([str(el).rjust(2) for el in self.lines[i]])}   {' '.join([str(el).rjust(2) for el in self.marks[i]])}\n"
        return outstr


def main():
    with open("aoc4-1.txt", "r") as file:
        lines = [line.strip() for line in file.readlines()]
    firstline = lines[0].split(",")
    lines = lines[2:]
    i = 0
    j = 0
    boards = []
    while i < len(lines):
        boards.append(Board(lines[i:i+5], boardnum=j))
        j += 1
        i += 6
    for el in firstline:
        print(el)
        is_to_pop = []
        for i, board in enumerate(boards):
            board.mark(el)
            if board.check():
                if len(boards) > 1:
                    is_to_pop.append(i)
                else:
                    return (board, int(el) * board.sum_unmarked())
        boards = [board for i, board in enumerate(boards) if i not in is_to_pop]

if __name__ == "__main__":
    board, total = main()
    print(board)
    print(total)