Files

118 lines
3.3 KiB
Python

from functools import reduce
from itertools import combinations
from textwrap import dedent
from typing import Set, override
from unittest import TestCase
from puzzles._solver import Solver
Point = tuple[int, int]
class DayEightSolver(Solver):
width: int
height: int
antenna_arrays: dict[str, set[Point]]
@override
def __init__(self, puzzle_input: str):
self.antenna_arrays = {}
lines = puzzle_input.strip().split("\n")
self.width = len(lines[0])
self.height = len(lines)
for x, line in enumerate(puzzle_input.strip().split("\n")):
for y, char in enumerate(line):
if char.isalnum():
if char not in self.antenna_arrays:
self.antenna_arrays[char] = set()
self.antenna_arrays[char].add((x, y))
@override
def solve_p1(self) -> int:
return len(
reduce(
Set.union,
(
self.antinodes(a, b, harmonic=False)
for antenna_array in self.antenna_arrays.values()
for a, b in combinations(antenna_array, 2)
),
)
)
@override
def solve_p2(self) -> int:
return len(
reduce(
Set.union,
(
self.antinodes(a, b, harmonic=True)
for antenna_array in self.antenna_arrays.values()
for a, b in combinations(antenna_array, 2)
),
)
)
def antinodes(self, a: Point, b: Point, harmonic=False) -> set[Point]:
points = set()
dx = b[0] - a[0]
dy = b[1] - a[1]
if harmonic:
antinode = a
while self.is_in_grid(antinode):
points.add(antinode)
antinode = (antinode[0] - dx, antinode[1] - dy)
antinode = b
while self.is_in_grid(antinode):
points.add(antinode)
antinode = (antinode[0] + dx, antinode[1] + dy)
else:
antinode = (a[0] - dx, a[1] - dy)
if self.is_in_grid(antinode):
points.add(antinode)
antinode = (b[0] + dx, b[1] + dy)
if self.is_in_grid(antinode):
points.add(antinode)
return points
def is_in_grid(self, point: Point) -> bool:
return 0 <= point[0] < self.height and 0 <= point[1] < self.width
class TestDayEightSolver(TestCase):
def test(self):
solver = DayEightSolver(
dedent(
"""
............
........0...
.....0......
.......0....
....0.......
......A.....
............
............
........A...
.........A..
............
............
"""
)
)
self.assertEqual(solver.width, 12)
self.assertEqual(solver.height, 12)
self.assertSetEqual(
solver.antenna_arrays["0"],
{(1, 8), (2, 5), (3, 7), (4, 4)},
)
self.assertEqual(solver.solve_p1(), 14)
self.assertEqual(solver.solve_p2(), 34)