Advent of Code 2021

day4.py

import collections

class Board:
    def __init__(self, rows):
        self.rows = [[int(n) for n in r.split()] for r in rows]
        self.all = collections.defaultdict(list)
        self.done = False

        for i, row in enumerate(self.rows):
            for j, n in enumerate(row):
                self.all[n].append((i, j))

    def is_win_row(self, row):
        return all(n == -1 for n in self.rows[row])

    def is_win_col(self, col):
        return all(row[col] == -1 for row in self.rows)

    def calculate_score(self):
        score = 0
        for row in self.rows:
            for n in row:
                score += n if n > -1 else 0

        return score

    def mark(self, n):
        if self.done:
            return -1

        if n in self.all:
            positions = self.all[n]

            for (i, j) in positions:
                self.rows[i][j] = -1

                if self.is_win_row(i) or self.is_win_col(j):
                    for (i, j) in positions:
                        self.rows[i][j] = -1

                    self.done = True
                    return self.calculate_score()

        return -1

with open("input") as f:
    lines = f.read().split("\n")

guesses = [int(n) for n in lines[0].split(",")]

boards = []
for i in range(2, len(lines), 6):
    boards.append(Board(lines[i : i + 5]))

def part1(guesses, boards):
    for guess in guesses:
        for b in boards:
            v = b.mark(guess)
            if v > -1:
                return v * guess

def part2(guesses, boards):
    scores = []

    for guess in guesses:
        for b in boards:
            v = b.mark(guess)
            if v > -1:
                scores.append(v * guess)

    return scores[-1]

print(part1(guesses, boards))
print(part2(guesses, boards))