aoc2021/4/b.py

113 lines
3.0 KiB
Python

import sys
WIDTH = 5
HEIGHT = 5
class Board(object):
def __init__(self, matrix, win_matrix):
# matrix: map of num => (x,y)
# win_matrix: map of (x,y) => picked bool
self.matrix = matrix
self.win_matrix = win_matrix
self.total = None # this is calculated upon noticing the board has won
@staticmethod
def from_block(block):
# convert the block, which is an array of string, into a flat list of numbers
nums = []
for row in block:
nums.extend([int(i) for i in row.strip().split()])
i = 0
matrix = {}
win_matrix = {}
for y in range(0, HEIGHT):
for x in range(0, WIDTH):
sq = nums[i]
matrix[sq] = (x, y, )
win_matrix[(x, y, )] = False
i += 1
return Board(matrix, win_matrix)
def fill(self, num):
# mark the space with $num as filled
# return whether this makes the board a winner
try:
filled_position = self.matrix[num]
except KeyError:
return False # number not on this board
self.win_matrix[filled_position] = True
return self.check_solved(num, *filled_position)
def check_solved(self, num, new_x, new_y):
# the square at ($new_x, $new_y) has been filled, check if it solves the board
# it solves the board if:
# the row is filled
# the col is filled
# either diagonal is filled ----- jk, the problem says diagonals dont matter
w = any([
# row
all([self.win_matrix[(x_coord, new_y, )] for x_coord in range(0, WIDTH)]),
# col
all([self.win_matrix[(new_x, y_coord, )] for y_coord in range(0, HEIGHT)])
])
if w:
self.total = self.calc_total(num)
return w
def calc_total(self, num):
print("calculating total")
# sum of all unmarked nums
unmarked = 0
for sq_num, coord in self.matrix.items():
if not self.win_matrix[coord]:
unmarked += sq_num
return unmarked * num
def update_boards(boards, pick):
winners = []
for board in boards:
if board.fill(pick):
winners.append(board)
return winners
with open("input.txt") as f:
lines = [i for i in f.readlines()]
# input parser
picks = [int(i) for i in lines.pop(0).strip().split(",")]
lines.pop(0)
board_blocks = []
while True:
# pop the 5 expected lines for the board
board_blocks.append([lines.pop(0), lines.pop(0), lines.pop(0), lines.pop(0), lines.pop(0), ])
if not lines:
break
empty = lines.pop(0)
boards = []
for block in board_blocks:
boards.append(Board.from_block(block))
for pick_num, pick in enumerate(picks):
winners = update_boards(boards, pick)
for winner in winners:
if len(boards) == 1:
print(boards[0].total)
sys.exit(0)
boards.remove(winner)
# 38m