139 lines
3.8 KiB
Python
139 lines
3.8 KiB
Python
from dataclasses import dataclass
|
|
from textwrap import dedent
|
|
from typing import NamedTuple, override
|
|
from unittest import TestCase
|
|
|
|
from puzzles._solver import Solver
|
|
|
|
|
|
class TrailSet(NamedTuple):
|
|
trailhead: tuple[int, int]
|
|
trailends: set[tuple[int, int]]
|
|
count: int
|
|
|
|
|
|
@dataclass
|
|
class TopographicalMap:
|
|
grid: list[list[str]]
|
|
trailheads: set[tuple[int, int]]
|
|
|
|
@classmethod
|
|
def parse(cls, raw_map: str) -> "TopographicalMap":
|
|
grid = []
|
|
trailheads = set()
|
|
|
|
for x, row in enumerate(raw_map.strip().splitlines()):
|
|
grid_row = []
|
|
|
|
for y, cell in enumerate(row):
|
|
height = int(cell)
|
|
grid_row.append(height)
|
|
if height == 0:
|
|
trailheads.add((x, y))
|
|
|
|
grid.append(grid_row)
|
|
|
|
return cls(grid, trailheads)
|
|
|
|
def trails(self) -> list[TrailSet]:
|
|
trails = []
|
|
|
|
for trailhead in self.trailheads:
|
|
trailends: set[tuple, tuple] = set()
|
|
count = 0
|
|
x, y = trailhead
|
|
height = self.grid[x][y]
|
|
|
|
def follow_trail(x, y, height):
|
|
nonlocal count
|
|
|
|
if height == 9:
|
|
trailends.add((x, y))
|
|
count += 1
|
|
return
|
|
|
|
if x > 0 and self.grid[x - 1][y] == height + 1:
|
|
follow_trail(x - 1, y, height + 1)
|
|
if x < len(self.grid) - 1 and self.grid[x + 1][y] == height + 1:
|
|
follow_trail(x + 1, y, height + 1)
|
|
if y > 0 and self.grid[x][y - 1] == height + 1:
|
|
follow_trail(x, y - 1, height + 1)
|
|
if y < len(self.grid[0]) - 1 and self.grid[x][y + 1] == height + 1:
|
|
follow_trail(x, y + 1, height + 1)
|
|
|
|
follow_trail(x, y, height)
|
|
|
|
trails.append(TrailSet(trailhead, trailends, count))
|
|
|
|
return trails
|
|
|
|
|
|
class DayTenSolver(Solver):
|
|
topographical_map: TopographicalMap
|
|
|
|
@override
|
|
def __init__(self, puzzle_input: str):
|
|
self.topographical_map = TopographicalMap.parse(puzzle_input)
|
|
|
|
@override
|
|
def solve_p1(self) -> int:
|
|
trail_sets = self.topographical_map.trails()
|
|
return sum(len(trail_set.trailends) for trail_set in trail_sets)
|
|
|
|
@override
|
|
def solve_p2(self) -> int:
|
|
trail_sets = self.topographical_map.trails()
|
|
return sum(trail_set.count for trail_set in trail_sets)
|
|
|
|
|
|
class TestDayTenSolver(TestCase):
|
|
test_input = dedent(
|
|
"""
|
|
89010123
|
|
78121874
|
|
87430965
|
|
96549874
|
|
45678903
|
|
32019012
|
|
01329801
|
|
10456732
|
|
"""
|
|
)
|
|
|
|
def test_parse(self):
|
|
solver = DayTenSolver(self.test_input)
|
|
self.assertEqual(
|
|
solver.topographical_map,
|
|
TopographicalMap(
|
|
grid=[
|
|
[8, 9, 0, 1, 0, 1, 2, 3],
|
|
[7, 8, 1, 2, 1, 8, 7, 4],
|
|
[8, 7, 4, 3, 0, 9, 6, 5],
|
|
[9, 6, 5, 4, 9, 8, 7, 4],
|
|
[4, 5, 6, 7, 8, 9, 0, 3],
|
|
[3, 2, 0, 1, 9, 0, 1, 2],
|
|
[0, 1, 3, 2, 9, 8, 0, 1],
|
|
[1, 0, 4, 5, 6, 7, 3, 2],
|
|
],
|
|
trailheads={
|
|
(0, 2),
|
|
(0, 4),
|
|
(2, 4),
|
|
(4, 6),
|
|
(5, 2),
|
|
(5, 5),
|
|
(6, 0),
|
|
(6, 6),
|
|
(7, 1),
|
|
},
|
|
),
|
|
)
|
|
|
|
def test_solve_p1(self):
|
|
solver = DayTenSolver(self.test_input)
|
|
self.assertEqual(solver.solve_p1(), 36)
|
|
|
|
def test_solve_p2(self):
|
|
solver = DayTenSolver(self.test_input)
|
|
self.assertEqual(solver.solve_p2(), 81)
|