96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
import itertools as it
|
|
from typing import Any, Iterator
|
|
|
|
import numpy as np
|
|
import parse # type: ignore
|
|
from numpy.typing import NDArray
|
|
|
|
from ..base import BaseSolver
|
|
|
|
|
|
class Solver(BaseSolver):
|
|
def part1(
|
|
self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], row: int
|
|
) -> int:
|
|
no_beacons_row_l: list[NDArray[np.floating[Any]]] = []
|
|
|
|
for (sx, sy), (bx, by) in sensor_to_beacon.items():
|
|
d = abs(sx - bx) + abs(sy - by) # closest
|
|
|
|
no_beacons_row_l.append(sx - np.arange(0, d - abs(sy - row) + 1)) # type: ignore
|
|
no_beacons_row_l.append(sx + np.arange(0, d - abs(sy - row) + 1)) # type: ignore
|
|
|
|
beacons_at_row = set(bx for (bx, by) in sensor_to_beacon.values() if by == row)
|
|
no_beacons_row = set(it.chain(*no_beacons_row_l)).difference(beacons_at_row) # type: ignore
|
|
|
|
return len(no_beacons_row)
|
|
|
|
def part2_intervals(
|
|
self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
|
|
) -> tuple[int, int, int]:
|
|
for y in self.progress.wrap(range(xy_max + 1)):
|
|
its: list[tuple[int, int]] = []
|
|
for (sx, sy), (bx, by) in sensor_to_beacon.items():
|
|
d = abs(sx - bx) + abs(sy - by)
|
|
dx = d - abs(sy - y)
|
|
|
|
if dx >= 0:
|
|
its.append((max(0, sx - dx), min(sx + dx, xy_max)))
|
|
|
|
its = sorted(its)
|
|
_, e = its[0]
|
|
|
|
for si, ei in its[1:]:
|
|
if si > e + 1:
|
|
return si - 1, y, 4_000_000 * (si - 1) + y
|
|
if ei > e:
|
|
e = ei
|
|
|
|
return (0, 0, 0)
|
|
|
|
def part2_cplex(
|
|
self, sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
|
|
) -> tuple[int, int, int]:
|
|
from docplex.mp.model import Model
|
|
|
|
m = Model()
|
|
|
|
x, y = m.continuous_var_list(2, ub=xy_max, name=["x", "y"])
|
|
|
|
for (sx, sy), (bx, by) in sensor_to_beacon.items():
|
|
d = abs(sx - bx) + abs(sy - by)
|
|
m.add_constraint(
|
|
m.abs(x - sx) + m.abs(y - sy) >= d + 1, # type: ignore
|
|
ctname=f"ct_{sx}_{sy}",
|
|
)
|
|
|
|
m.set_objective("min", x + y)
|
|
|
|
s = m.solve()
|
|
assert s is not None
|
|
|
|
vx = int(s.get_value(x))
|
|
vy = int(s.get_value(y))
|
|
return vx, vy, 4_000_000 * vx + vy
|
|
|
|
def solve(self, input: str) -> Iterator[Any]:
|
|
lines = input.splitlines()
|
|
|
|
sensor_to_beacon: dict[tuple[int, int], tuple[int, int]] = {}
|
|
|
|
for line in lines:
|
|
r: dict[str, str] = parse.parse( # type: ignore
|
|
"Sensor at x={sx}, y={sy}: closest beacon is at x={bx}, y={by}", line
|
|
)
|
|
sensor_to_beacon[int(r["sx"]), int(r["sy"])] = (int(r["bx"]), int(r["by"]))
|
|
|
|
xy_max = 4_000_000 if max(sensor_to_beacon) > (1_000, 0) else 20
|
|
row = 2_000_000 if max(sensor_to_beacon) > (1_000, 0) else 10
|
|
|
|
yield self.part1(sensor_to_beacon, row)
|
|
|
|
# x, y, a2 = part2_cplex(sensor_to_beacon, xy_max)
|
|
x, y, a2 = self.part2_intervals(sensor_to_beacon, xy_max)
|
|
self.logger.info(f"answer 2 is {a2} (x={x}, y={y})")
|
|
yield a2
|