AOC 2021 Day 4

day04.py

#!/usr/bin/python3
import numpy as np

DATA = 'data/day04'

def total(board):
    return np.sum(board[board > 0])

def main():
    boards = []
    board = []

    with open(DATA) as fp:
        numbers = [int(n) for n in fp.readline().strip().split(',')]

        for line in fp:
            line = line.strip()
            if line:
                board.append([int(n) for n in line.split()])
            elif board:
                boards.append(board)
                board = []

    if board:
        boards.append(board)

    boards = np.array(boards)
    tot_boards = len(boards)

    for num in numbers:
        boards[boards == num] = -1

        winners = [i for i, b in enumerate(boards)
                if np.any(np.all(b < 0, axis=0))
                or np.any(np.all(b < 0, axis=1))]

        if winners:
            if len(boards) == tot_boards:
                print('P1 =', total(boards[winners[0]]) * num)

            if len(boards) == len(winners):
                print('P2 =', total(boards[winners[-1]]) * num)
                break

            boards = np.delete(boards, winners, 0)

if __name__ == "__main__":
    main()

Ohh, and numpy deals with the dimensions when doing that internally too? Neat! Thanks, I'll definitely have to dive a bit deeper into numpy as these challenges get more complicated.