adventofcode 2021 day 04

adventofcode_2021_04.py

#!/usr/bin/env python3

from pprint import pprint
from typing import List, NamedTuple, Optional

with open('data.txt') as fp:
    lines = iter(fp.read().splitlines(False))


class Num(int):
    drawn: bool = False


class Color:
    RED = '\x1b[31m'
    GREEN = '\x1b[32m'
    WHITE = '\x1b[37m'
    RESET = '\x1b[0m'


numbers = map(Num, next(lines).split(','))

boards: List[List[List[Num]]] = []

for line in lines:
    if not line.strip():
        boards.append([])
        continue

    boards[-1].append(list(map(Num, filter(str.strip, line.split(' ')))))

pprint(boards, indent=2)


def render_board(board, winner: bool = False):
    for idx, line in enumerate(board):
        for num in line:
            print(f"{Color.RED if num.drawn else Color.WHITE}{num:02d}{Color.RESET} ", end='')
        if winner and not idx:
            print(f" - {Color.GREEN}WINNER{Color.RESET}", end='')
        print()
    print()


def columns(board):
    return list(map(list, zip(*board)))


def is_completed(board) -> bool:
    # check rows
    if any(all(num.drawn for num in row) for row in board):
        return True

    # check cols
    if any(all(num.drawn for num in col) for col in columns(board)):
        return True

    return False


winners = []

for draw, number in enumerate(numbers, start=1):
    input('draw next number? ')
    print('-' * 3, f"{number:02d}", '-' * 3)

    for idx, board in enumerate(boards, start=1):
        if idx not in [winner[1] for winner in winners]:

            for row in board:
                if number in row:
                    row[row.index(number)].drawn = True

            if is_completed(board):
                winners.append((board, idx, number, draw))

        render_board(board, idx in [winner[1] for winner in winners])

    ## uncommented to cover part2
    # if winners:
    #     break

    if len(winners) == len(boards):
        print('all boards have won')
        break

print(f'got {len(winners)} winner(s):')

for board, idx, number, draw in winners:
    usum = sum(n for r in board for n in r if not n.drawn)
    print(f" - board={idx} {draw=:02d} {number=:02d} usum={usum} result={usum * number}")

print('DONE')

data.txt

7,4,9,5,11,17,23,2,0,14,21,24,10,16,13,6,15,25,12,22,18,20,8,19,3,26,1

22 13 17 11  0
 8  2 23  4 24
21  9 14 16  7
 6 10  3 18  5
 1 12 20 15 19

 3 15  0  2 22
 9 18 13 17  5
19  8  7 25 23
20 11 10 24  4
14 21 16 12  6

14 21 17 24  4
10 16 15  9 19
18  8 23 26 20
22 11 13  6  5
 2  0 12  3  7