81 lines
2.2 KiB
Python
81 lines
2.2 KiB
Python
|
#!/usr/bin/env python
|
||
|
|
|
||
|
|
import dataclasses
|
||
|
|
import functools
|
||
|
|
import itertools
|
||
|
|
import sys
|
||
|
|
from typing import Literal, NamedTuple
|
||
|
|
|
||
|
|
|
||
|
|
class Point(NamedTuple):
|
||
|
|
x: int
|
||
|
|
y: int
|
||
|
|
|
||
|
|
|
||
|
|
@dataclasses.dataclass
|
||
|
|
class Robot:
|
||
|
|
pos: Point
|
||
|
|
vel: Point
|
||
|
|
|
||
|
|
def step(self, dims: Point, delta: int = 1) -> "Robot":
|
||
|
|
x, y = self.pos.x + self.vel.x * delta, self.pos.y + self.vel.y * delta
|
||
|
|
return Robot(
|
||
|
|
Point(x % dims.x, y % dims.y),
|
||
|
|
self.vel,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def solve(input: str) -> int:
|
||
|
|
def parse_robot(input: str) -> Robot:
|
||
|
|
pos, vel = map(lambda s: s.split("=")[1], input.split(" "))
|
||
|
|
return Robot(
|
||
|
|
Point(*map(int, pos.split(","))),
|
||
|
|
Point(*map(int, vel.split(","))),
|
||
|
|
)
|
||
|
|
|
||
|
|
def parse(input: list[str]) -> list[Robot]:
|
||
|
|
return [parse_robot(line) for line in input]
|
||
|
|
|
||
|
|
def find_tree(robots: list[Robot], dims: Point) -> int:
|
||
|
|
def compute_positions(step: int) -> list[Point]:
|
||
|
|
return [robot.step(dims, step).pos for robot in robots]
|
||
|
|
|
||
|
|
def compute_variance(values: list[int]) -> float:
|
||
|
|
avg = sum(values) / len(values)
|
||
|
|
variance = sum((n - avg) ** 2 for n in values) / len(values)
|
||
|
|
return variance
|
||
|
|
|
||
|
|
def cluster_variance(step: int, dimension: Literal["x", "y"]) -> float:
|
||
|
|
return compute_variance(
|
||
|
|
[getattr(p, dimension) for p in compute_positions(step)]
|
||
|
|
)
|
||
|
|
|
||
|
|
# The tree should have robots clustered together in X and Y
|
||
|
|
cluster_x = min(
|
||
|
|
range(dims.x),
|
||
|
|
key=functools.partial(cluster_variance, dimension="x"),
|
||
|
|
)
|
||
|
|
cluster_y = min(
|
||
|
|
range(dims.y),
|
||
|
|
key=functools.partial(cluster_variance, dimension="y"),
|
||
|
|
)
|
||
|
|
|
||
|
|
# And those clusers should repeat modulo each dimension
|
||
|
|
for i in itertools.count(cluster_x, step=dims.x):
|
||
|
|
if i % dims.y == cluster_y:
|
||
|
|
return i
|
||
|
|
assert False # Sanity check
|
||
|
|
|
||
|
|
robots = parse(input.splitlines())
|
||
|
|
dims = Point(101, 103)
|
||
|
|
return find_tree(robots, dims)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
input = sys.stdin.read()
|
||
|
|
print(solve(input))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|