134 lines
3.6 KiB
Python
134 lines
3.6 KiB
Python
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)
|