Start adding progress.

src/holt59/aoc/2022/day15.py

@@ -1,90 +1,96 @@
 import sys
-from typing import Any
+from typing import Any, Iterator
 
 import numpy as np
 import parse  # type: ignore
 from numpy.typing import NDArray
 
-
-def part1(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(np.concatenate(no_beacons_row_l)).difference(beacons_at_row)  # type: ignore
-
-    return len(no_beacons_row)
+from ..base import BaseSolver
 
 
-def part2_intervals(
-    sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
-) -> tuple[int, int, int]:
-    from tqdm import trange
+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(np.concatenate(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 y in trange(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)
+            m.add_constraint(
+                m.abs(x - sx) + m.abs(y - sy) >= d + 1, ctname=f"ct_{sx}_{sy}"
+            )  # type: ignore
 
-            if dx >= 0:
-                its.append((max(0, sx - dx), min(sx + dx, xy_max)))
+        m.set_objective("min", x + y)
 
-        its = sorted(its)
-        _, e = its[0]
+        s = m.solve()
+        assert s is not None
 
-        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
+        vx = int(s.get_value(x))
+        vy = int(s.get_value(y))
+        return vx, vy, 4_000_000 * vx + vy
 
-    return (0, 0, 0)
+    def solve(self, input: str) -> Iterator[Any]:
+        lines = input.splitlines()
 
+        sensor_to_beacon: dict[tuple[int, int], tuple[int, int]] = {}
 
-def part2_cplex(
-    sensor_to_beacon: dict[tuple[int, int], tuple[int, int]], xy_max: int
-) -> tuple[int, int, int]:
-    from docplex.mp.model import Model
+        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"]))
 
-    m = Model()
+        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
 
-    x, y = m.continuous_var_list(2, ub=xy_max, name=["x", "y"])
+        yield self.part1(sensor_to_beacon, row)
 
-    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, ctname=f"ct_{sx}_{sy}")  # type: ignore
-
-    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
-
-
-lines = sys.stdin.read().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
-
-print(f"answer 1 is {part1(sensor_to_beacon, row)}")
-
-# x, y, a2 = part2_cplex(sensor_to_beacon, xy_max)
-x, y, a2 = part2_intervals(sensor_to_beacon, xy_max)
-print(f"answer 2 is {a2} (x={x}, y={y})")
+        # x, y, a2 = part2_cplex(sensor_to_beacon, xy_max)
+        x, y, a2 = self.part2_intervals(sensor_to_beacon, xy_max)
+        self.logger.info("answer 2 is {at} (x={x}, y={y})")
+        yield a2