2021/day04.py

87 lines
1.9 KiB
Python
Raw Normal View History

2021-12-09 11:07:49 +00:00
import fileinput
from typing import List, Tuple, Set
numbers = [int(x) for x in input().split(",")]
input()
Board = List[List[int]]
boards: List[Board] = []
board: Board = []
while True:
try:
line = input().strip()
except EOFError:
break
if line:
board.append([int(x) for x in line.split()])
else:
boards.append(board)
board = []
if board:
boards.append(board)
def find_bingo(drawn: Set[int], wins: Set[int]) -> Tuple[int, int]:
for n in numbers:
drawn.add(n)
for b, board in enumerate(boards):
for row in board:
for col in row:
if col not in drawn:
break
else:
if b not in wins:
wins.add(b)
return b, n
for x in range(len(board[0])):
for y in range(len(board)):
if board[y][x] not in drawn:
break
else:
if b not in wins:
wins.add(b)
return b, n
raise RuntimeError("no bingo found")
def print_board(board, drawn=frozenset()):
for r in board:
for c in r:
if c in drawn:
print("\033[91m", end="")
print(str(c).rjust(4), end="\033[0m")
print()
print()
print(", ".join(str(n) for n in numbers))
print()
for b in boards:
print_board(b)
print("===")
drawn = set()
wins = set()
b, last_drawn = find_bingo(drawn, wins)
undrawn = sum(number for row in boards[b] for number in row if number not in drawn)
part1 = last_drawn * undrawn
while len(wins) < len(boards):
b, last_drawn = find_bingo(drawn, wins)
print_board(boards[b], drawn)
undrawn = sum(number for row in boards[b] for number in row if number not in drawn)
part2 = last_drawn * undrawn
print("part1:", part1)
print("part2:", part2)