Restructure solver modules with classes and unittest suites
This commit is contained in:
155
puzzles/4.py
155
puzzles/4.py
@@ -1,83 +1,90 @@
|
||||
from typing import Iterator
|
||||
from textwrap import dedent
|
||||
from typing import Iterator, override
|
||||
from unittest import TestCase
|
||||
|
||||
test_input = """
|
||||
MMMSXXMASM
|
||||
MSAMXMSMSA
|
||||
AMXSXMAAMM
|
||||
MSAMASMSMX
|
||||
XMASAMXAMM
|
||||
XXAMMXXAMA
|
||||
SMSMSASXSS
|
||||
SAXAMASAAA
|
||||
MAMMMXMMMM
|
||||
MXMXAXMASX
|
||||
"""
|
||||
from more_itertools import sliding_window
|
||||
|
||||
test_solution_p1 = 18
|
||||
test_solution_p2 = 9
|
||||
from puzzles._solver import Solver
|
||||
|
||||
Line = tuple[str, ...]
|
||||
Grid = tuple[Line, ...]
|
||||
|
||||
|
||||
def solve_p1(puzzle_input: str) -> int:
|
||||
word_grid = _parse_world_grid(puzzle_input)
|
||||
return sum(_count_xmas(line) for line in _scan_word_grid(word_grid))
|
||||
class DayFourSolver(Solver):
|
||||
grid: Grid
|
||||
|
||||
@override
|
||||
def __init__(self, puzzle_input: str):
|
||||
self.grid = tuple(tuple(row) for row in puzzle_input.strip().splitlines())
|
||||
|
||||
def solve_p2(puzzle_input: str) -> int:
|
||||
word_grid = _parse_world_grid(puzzle_input)
|
||||
squares = _three_squares(word_grid)
|
||||
return sum(1 for square in squares if _has_cross_mas(square))
|
||||
|
||||
|
||||
def _parse_world_grid(puzzle_input: str) -> tuple[tuple[str, ...], ...]:
|
||||
return tuple(tuple(row) for row in puzzle_input.strip().splitlines())
|
||||
|
||||
|
||||
def _scan_word_grid(
|
||||
char_grid: tuple[tuple[str, ...], ...]
|
||||
) -> Iterator[tuple[str, ...]]:
|
||||
yield from (row for row in char_grid)
|
||||
yield from (tuple(col) for col in zip(*char_grid))
|
||||
yield from (
|
||||
_diagonal(char_grid, i) for i in range(-len(char_grid) + 1, len(char_grid[0]))
|
||||
)
|
||||
yield from (
|
||||
_diagonal(tuple(reversed(char_grid)), i)
|
||||
for i in range(-len(char_grid) + 1, len(char_grid[0]))
|
||||
)
|
||||
|
||||
|
||||
def _diagonal(grid: tuple[tuple[str, ...], ...], offset=0) -> tuple[str, ...]:
|
||||
return tuple(
|
||||
grid[i][i + offset] for i in range(len(grid)) if 0 <= i + offset < len(grid)
|
||||
)
|
||||
|
||||
|
||||
def _count_xmas(line: tuple[str, ...]) -> int:
|
||||
return sum(
|
||||
1
|
||||
for i in range(len(line) - 3)
|
||||
if line[i : i + 4] in (("X", "M", "A", "S"), ("S", "A", "M", "X"))
|
||||
)
|
||||
|
||||
|
||||
ThreeSquare = tuple[tuple[str, str, str], tuple[str, str, str], tuple[str, str, str]]
|
||||
|
||||
|
||||
def _three_squares(word_grid: tuple[tuple[str, ...], ...]) -> Iterator[ThreeSquare]:
|
||||
yield from (
|
||||
(
|
||||
word_grid[i][j : j + 3],
|
||||
word_grid[i + 1][j : j + 3],
|
||||
word_grid[i + 2][j : j + 3],
|
||||
@override
|
||||
def solve_p1(self) -> int:
|
||||
targets = (("X", "M", "A", "S"), ("S", "A", "M", "X"))
|
||||
return sum(
|
||||
1
|
||||
for line in self.scan_lines()
|
||||
for window in sliding_window(line, 4)
|
||||
if window in targets
|
||||
)
|
||||
|
||||
@override
|
||||
def solve_p2(self) -> int:
|
||||
targets = (("M", "A", "S"), ("S", "A", "M"))
|
||||
return sum(
|
||||
1
|
||||
for square in self.scan_squares(3)
|
||||
if (square[0][0], square[1][1], square[2][2]) in targets
|
||||
and (square[0][2], square[1][1], square[2][0]) in targets
|
||||
)
|
||||
|
||||
def scan_lines(self) -> Iterator[Line]:
|
||||
yield from self.grid
|
||||
for col in zip(*self.grid):
|
||||
yield tuple(col)
|
||||
for i in range(-len(self.grid) + 1, len(self.grid[0])):
|
||||
yield self.diagonal(i)
|
||||
yield self.diagonal(i, inverse=True)
|
||||
|
||||
def scan_squares(self, size=3) -> Iterator[Grid]:
|
||||
yield from (
|
||||
(
|
||||
self.grid[i][j : j + size],
|
||||
self.grid[i + 1][j : j + size],
|
||||
self.grid[i + 2][j : j + size],
|
||||
)
|
||||
for i in range(len(self.grid) - size + 1)
|
||||
for j in range(len(self.grid[0]) - size + 1)
|
||||
)
|
||||
|
||||
def diagonal(self, offset=0, inverse=False) -> Line:
|
||||
return tuple(
|
||||
self.grid[i][len(self.grid[0]) - i - offset - 1 if inverse else i + offset]
|
||||
for i in range(len(self.grid))
|
||||
if 0 <= i + offset < len(self.grid)
|
||||
)
|
||||
for i in range(len(word_grid) - 2)
|
||||
for j in range(len(word_grid[0]) - 2)
|
||||
)
|
||||
|
||||
|
||||
def _has_cross_mas(square: ThreeSquare):
|
||||
diag_1 = (square[0][0], square[1][1], square[2][2])
|
||||
diag_2 = (square[0][2], square[1][1], square[2][0])
|
||||
return (diag_1 == ("M", "A", "S") or diag_1 == ("S", "A", "M")) and (
|
||||
diag_2 == ("M", "A", "S") or diag_2 == ("S", "A", "M")
|
||||
)
|
||||
class TestDayFourSolver(TestCase):
|
||||
def test(self):
|
||||
solver = DayFourSolver(
|
||||
dedent(
|
||||
"""
|
||||
MMMSXXMASM
|
||||
MSAMXMSMSA
|
||||
AMXSXMAAMM
|
||||
MSAMASMSMX
|
||||
XMASAMXAMM
|
||||
XXAMMXXAMA
|
||||
SMSMSASXSS
|
||||
SAXAMASAAA
|
||||
MAMMMXMMMM
|
||||
MXMXAXMASX
|
||||
"""
|
||||
)
|
||||
)
|
||||
self.assertEqual(
|
||||
solver.grid[0],
|
||||
("M", "M", "M", "S", "X", "X", "M", "A", "S", "M"),
|
||||
)
|
||||
self.assertEqual(solver.solve_p1(), 18)
|
||||
self.assertEqual(solver.solve_p2(), 9)
|
||||
|
||||
Reference in New Issue
Block a user