Home — Advent of Code

import numpy as np

class Board():    
    def __init__(self, size):
        self.numbers = np.zeros((size, size), np.int32)
        self.marks = np.zeros((size, size), bool)
        
    def is_winner(self):
        size = self.marks.shape[0]
        for i in range(size):
            if np.sum(self.marks[i]) == size or \
                np.sum(self.marks[:,i]) == size:
                
                return True
    
    def mark(self, n):
        self.marks |= self.numbers == n
    
    def sum_unmarked(self):
        return np.sum(self.numbers * ~self.marks)

def parse_file(f):
    numbers_drawn = map(int, next(f).split(','))
    board = None
    boards = []
    
    i = 0
    
    for l in f:
        l = l.rstrip()
        if l == "":
            if board:
                boards.append(board)
            board = Board(5)
            i = 0
            continue
            
        board.numbers[i] = list(map(int, l.split()))
        i += 1
        
    return numbers_drawn, boards

def find_winner(numbers_drawn, boards):
    first_winner_score = None
    last_winner_score = None
    numbers_drawn = list(numbers_drawn)
    for ni, n in enumerate(numbers_drawn):
        for b in boards:
            b.mark(n)
            
        for bi, b in enumerate(boards):
            if b.is_winner():
                boards[bi] = None
                
                if first_winner_score is None:
                    first_winner_score = b.sum_unmarked() * n
                    
                if len(boards) == 1:
                    last_winner_score  = b.sum_unmarked() * n

        boards = [b for b in boards if b]
                    
    return first_winner_score, last_winner_score

import import_ipynb
import helper
f = helper.open_file('2021_4.txt')
numbers_drawn, boards = parse_file(f)

find_winner(numbers_drawn, boards)