109 lines
3.3 KiB
Python
109 lines
3.3 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
import sys
|
||
|
|
from typing import Optional, Sequence
|
||
|
|
|
||
|
|
|
||
|
|
class Board:
|
||
|
|
def __init__(self, matrix: Sequence[Sequence[int]]):
|
||
|
|
self.matrix = matrix
|
||
|
|
self.marks = [[False] * len(matrix[0]) for _ in range(len(matrix[0]))]
|
||
|
|
|
||
|
|
def mark(self, num: int):
|
||
|
|
for i, row in enumerate(self.matrix):
|
||
|
|
for j, col in enumerate(row):
|
||
|
|
if col == num:
|
||
|
|
self.marks[i][j] = True
|
||
|
|
|
||
|
|
def row(self, i: int) -> Sequence[int]:
|
||
|
|
return self.matrix[i]
|
||
|
|
|
||
|
|
def col(self, j: int) -> Sequence[int]:
|
||
|
|
return [self.matrix[i][j] for i in range(len(self.matrix))]
|
||
|
|
|
||
|
|
def row_marks(self, i: int) -> Sequence[int]:
|
||
|
|
return self.marks[i]
|
||
|
|
|
||
|
|
def col_marks(self, j: int) -> Sequence[int]:
|
||
|
|
return [self.marks[i][j] for i in range(len(self.marks))]
|
||
|
|
|
||
|
|
def row_marked(self) -> Optional[int]:
|
||
|
|
"If an entire row is marked, return that row's index. Else return None."
|
||
|
|
for i in range(len(self.matrix)):
|
||
|
|
if all(self.row_marks(i)):
|
||
|
|
return i
|
||
|
|
return None
|
||
|
|
|
||
|
|
def col_marked(self) -> Optional[int]:
|
||
|
|
"If an entire column is marked, return that column's index. Else return None."
|
||
|
|
for j in range(len(self.matrix[0])):
|
||
|
|
if all(self.col_marks(j)):
|
||
|
|
return j
|
||
|
|
return None
|
||
|
|
|
||
|
|
def all_unmarked(self) -> Sequence[int]:
|
||
|
|
"Gets a list of all the unmarked numbers on this board"
|
||
|
|
return [
|
||
|
|
self.matrix[i][j]
|
||
|
|
for i in range(len(self.matrix))
|
||
|
|
for j in range(len(self.matrix[i]))
|
||
|
|
if not self.marks[i][j]
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def part1(draws: Sequence[int], matrices: Sequence[Sequence[Sequence[int]]]):
|
||
|
|
boards = [Board(matrix) for matrix in matrices]
|
||
|
|
result = None
|
||
|
|
for draw in draws:
|
||
|
|
# Mark the boards
|
||
|
|
for board in boards:
|
||
|
|
board.mark(draw)
|
||
|
|
# Check to see if there are 5 in a row in any direction
|
||
|
|
row_num = board.row_marked()
|
||
|
|
if row_num is not None:
|
||
|
|
result = sum(board.all_unmarked()) * draw
|
||
|
|
break
|
||
|
|
col_num = board.col_marked()
|
||
|
|
if col_num is not None:
|
||
|
|
result = sum(board.all_unmarked()) * draw
|
||
|
|
break
|
||
|
|
|
||
|
|
if result is not None:
|
||
|
|
break
|
||
|
|
print(result)
|
||
|
|
|
||
|
|
|
||
|
|
def part2(draws: Sequence[int], matrices: Sequence[Sequence[Sequence[int]]]):
|
||
|
|
boards = [Board(matrix) for matrix in matrices]
|
||
|
|
result = None
|
||
|
|
for draw in draws:
|
||
|
|
# Mark the boards
|
||
|
|
for board in boards:
|
||
|
|
board.mark(draw)
|
||
|
|
# Filter out the winners
|
||
|
|
if len(boards) == 1:
|
||
|
|
board = boards[0]
|
||
|
|
if board.row_marked() is not None or board.col_marked() is not None:
|
||
|
|
result = sum(board.all_unmarked()) * draw
|
||
|
|
break
|
||
|
|
else:
|
||
|
|
boards = [
|
||
|
|
board
|
||
|
|
for board in boards
|
||
|
|
if board.row_marked() is None and board.col_marked() is None
|
||
|
|
]
|
||
|
|
print(result)
|
||
|
|
|
||
|
|
|
||
|
|
lines = [line.strip() for line in sys.stdin if line.strip()]
|
||
|
|
draws = [int(draw) for draw in lines.pop(0).split(",")]
|
||
|
|
# Matrices are 5 lines
|
||
|
|
matrices = [
|
||
|
|
[[int(col) for col in line.split()] for line in lines[i : i + 5]]
|
||
|
|
for i in range(0, len(lines), 5)
|
||
|
|
]
|
||
|
|
|
||
|
|
print("Part 1")
|
||
|
|
part1(draws, matrices)
|
||
|
|
print("Part 2")
|
||
|
|
part2(draws, matrices)
|