diff --git a/puzzles/13.py b/puzzles/13.py index ade66bc..1d28993 100644 --- a/puzzles/13.py +++ b/puzzles/13.py @@ -2,6 +2,7 @@ import re from textwrap import dedent from typing import Iterator, NamedTuple, override from unittest import TestCase + from puzzles._solver import Solver diff --git a/puzzles/14.py b/puzzles/14.py new file mode 100644 index 0000000..ae0258d --- /dev/null +++ b/puzzles/14.py @@ -0,0 +1,133 @@ +import re +from collections import Counter +from itertools import count +from math import prod +from textwrap import dedent +from typing import NamedTuple, override +from unittest import TestCase + +from more_itertools import filter_map, first_true, iterate + +from puzzles._solver import Solver + + +class Robot(NamedTuple): + x: int + y: int + dx: int + dy: int + + +robot_spec_pattern = re.compile(r"p=(-?\d+),(-?\d+) v=(-?\d+),(-?\d+)") + + +class DayFourteenSolver(Solver): + robots: list[Robot] + width: int + height: int + + @override + def __init__(self, puzzle_input: str, width=101, height=103): + self.robots = [ + Robot(*map(int, match.groups())) + for match in robot_spec_pattern.finditer(puzzle_input) + ] + self.width = width + self.height = height + + @override + def solve_p1(self) -> int: + quadrant_counts = Counter( + filter_map( + lambda r: self.get_quadrant(self.move_robot(r, 100)), + self.robots, + ) + ) + return prod(quadrant_counts.values()) + + @override + def solve_p2(self) -> int: + state = iterate( + lambda robots: [self.move_robot(robot, 1) for robot in robots], + self.robots, + ) + return first_true( + count(), + pred=lambda i: self.total_overlap(next(state)) == 0, + ) + + def move_robot(self, robot: Robot, seconds: int) -> Robot: + return Robot( + (robot.x + robot.dx * seconds) % self.width, + (robot.y + robot.dy * seconds) % self.height, + robot.dx, + robot.dy, + ) + + def get_quadrant(self, robot: Robot) -> int | None: + if robot.x > self.width // 2 and robot.y < self.height // 2: + return 1 + if robot.x < self.width // 2 and robot.y < self.height // 2: + return 2 + if robot.x < self.width // 2 and robot.y > self.height // 2: + return 3 + if robot.x > self.width // 2 and robot.y > self.height // 2: + return 4 + return None + + def total_overlap(self, robots: list[Robot]) -> int: + locations = Counter((robot.x, robot.y) for robot in robots) + return sum(count for count in locations.values() if count > 1) + + +class TestDayFourteenSolver(TestCase): + def test(self): + solver = DayFourteenSolver( + dedent( + """ + p=0,4 v=3,-3 + p=6,3 v=-1,-3 + p=10,3 v=-1,2 + p=2,0 v=2,-1 + p=0,0 v=1,3 + p=3,0 v=-2,-2 + p=7,6 v=-1,-3 + p=3,0 v=-1,-2 + p=9,3 v=2,3 + p=7,3 v=-1,2 + p=2,4 v=2,-3 + p=9,5 v=-3,-3 + """ + ), + width=11, + height=7, + ) + self.assertTupleEqual( + solver.robots[0], + Robot(0, 4, 3, -3), + ) + self.assertTupleEqual( + solver.move_robot(solver.robots[0], 1), + Robot(3, 1, 3, -3), + ) + self.assertEqual(solver.get_quadrant(Robot(0, 0, 0, 0)), 2) + self.assertEqual(solver.get_quadrant(Robot(5, 0, 0, 0)), None) + self.assertEqual( + solver.total_overlap( + [ + Robot(0, 0, 0, 0), + Robot(0, 0, 0, 0), + ] + ), + 2, + ) + self.assertEqual( + solver.total_overlap( + [ + Robot(0, 0, 0, 0), + Robot(1, 0, 0, 0), + ] + ), + 0, + ) + self.assertEqual(solver.solve_p1(), 12) diff --git a/puzzles/9.py b/puzzles/9.py index 66b041e..3e50866 100644 --- a/puzzles/9.py +++ b/puzzles/9.py @@ -1,5 +1,5 @@ -from itertools import count, pairwise import math +from itertools import count, pairwise from typing import Iterable, NamedTuple, override from unittest import TestCase