AoC 2025 Day 8 - faster

Author: pareronia

src/main/python/AoC2025_08.py

@@ -3,17 +3,12 @@
 # Advent of Code 2025 Day 8
 #
 
-import itertools
 import sys
-from collections import defaultdict
-from collections import deque
-from collections.abc import Callable
-from collections.abc import Iterator
-from functools import cmp_to_key
 from math import prod
 
 from aoc.common import InputData
 from aoc.common import SolutionBase
+from aoc.geometry3d import Position3D
 
 TEST = """\
 162,817,812
@@ -38,113 +33,75 @@
 425,690,689
 """
 
-Box = tuple[int, int, int]
-Input = tuple[int, list[tuple[Box, Box]]]
+Input = tuple[list[Position3D], list[tuple[int, int, int]]]
 Output1 = int
 Output2 = int
 
 
-def connected_components(
-    nodes: set[Box], adjacent: Callable[[Box], Iterator[Box]]
-) -> list[set[Box]]:
-    components = list[set[Box]]()
-    todo = set(nodes)
-    while len(todo) != 0:
-        node = todo.pop()
-        component = {node}
-        q = deque[Box]([node])
-        while len(q) != 0:
-            node = q.popleft()
-            for n in adjacent(node):
-                if n not in component:
-                    q.append(n)
-                component.add(n)
-        components.append(component)
-    return components
+class DSU:
+    def __init__(self, size: int) -> None:
+        self.ids = list(range(size))
+        self.sz = [1 for _ in range(size)]
+        self.num_components = size
+
+    def find(self, p: int) -> int:
+        if self.ids[p] != p:
+            self.ids[p] = self.find(self.ids[p])
+        return self.ids[p]
+
+    def unify(self, p: int, q: int) -> None:
+        root_p, root_q = self.find(p), self.find(q)
+        if root_p != root_q:
+            if self.sz[root_p] < self.sz[root_q]:
+                self.sz[root_q] += self.sz[root_p]
+                self.ids[root_p] = root_q
+                self.sz[root_p] = 0
+            else:
+                self.sz[root_p] += self.sz[root_q]
+                self.ids[root_q] = root_p
+                self.sz[root_q] = 0
+            self.num_components -= 1
 
 
 class Solution(SolutionBase[Input, Output1, Output2]):
     def parse_input(self, input_data: InputData) -> Input:
-        def cmp(p1: tuple[Box, Box], p2: tuple[Box, Box]) -> int:
-            def distance(a: Box, b: Box) -> float:
-                x1, y1, z1 = a
-                x2, y2, z2 = b
-                dx, dy, dz = x1 - x2, y1 - y2, z1 - z2
-                return dx * dx + dy * dy + dz * dz
-
-            dp1 = distance(p1[0], p1[1])
-            dp2 = distance(p2[0], p2[1])
-            if dp1 > dp2:
-                return 1
-            if dp1 < dp2:
-                return -1
-            return 0
-
-        boxes = list[Box]()
-        for line in input_data:
-            x, y, z = map(int, line.split(","))
-            boxes.append((x, y, z))
-        return len(boxes), sorted(
-            itertools.combinations(boxes, 2), key=cmp_to_key(cmp)
+        boxes = [
+            Position3D.of(int(x), int(y), int(z))
+            for x, y, z in (line.split(",") for line in input_data)
+        ]
+        pairs = sorted(
+            (
+                (i, j, boxes[i].squared_distance(boxes[j]))
+                for i in range(len(boxes))
+                for j in range(i + 1, len(boxes))
+            ),
+            key=lambda pair: pair[2],
         )
+        return boxes, pairs
 
-    def get_components(
-        self, edges: dict[Box, set[Box]]
-    ) -> set[frozenset[Box]]:
-        return {
-            frozenset(c)
-            for c in connected_components(
-                set(edges.keys()), lambda n: (nxt for nxt in edges[n])
-            )
-        }
-
-    def solve_1(self, inputs: Input, num_pairs: int) -> int:
-        _, pairs = inputs
-        edges = defaultdict[Box, set[Box]](set)
-        for i in range(num_pairs):
-            a, b = pairs[i]
-            edges[a].add(b)
-            edges[b].add(a)
-        components = self.get_components(edges)
-        best = sorted(components, key=len, reverse=True)
-        return prod(len(c) for c in best[:3])
+    def solve_1(self, inputs: Input, num_pairs: int = 1000) -> int:
+        boxes, pairs = inputs
+        dsu = DSU(len(boxes))
+        for pair in pairs[:num_pairs]:
+            dsu.unify(pair[0], pair[1])
+        return prod(sorted(dsu.sz)[-3:])
 
     def part_1(self, pairs: Input) -> Output1:
-        return self.solve_1(pairs, 1000)
-
-    def solve_2(self, inputs: Input, start: int) -> int:
-        num_boxes, pairs = inputs
-        edges = defaultdict[Box, set[Box]](set)
-        edges = defaultdict[Box, set[Box]](set)
-        for i in range(start):
-            a, b = pairs[i]
-            edges[a].add(b)
-            edges[b].add(a)
-        components = self.get_components(edges)
-        while i < len(pairs) - 1:
-            i += 1
-            a, b = pairs[i]
-            for c in components:
-                if a in c and b in c:
-                    break
-            else:
-                edges[a].add(b)
-                edges[b].add(a)
-                components = self.get_components(edges)
-                if (
-                    len(components) == 1
-                    and len(next(_ for _ in components)) == num_boxes
-                ):
-                    return pairs[i][0][0] * pairs[i][1][0]
+        return self.solve_1(pairs)
+
+    def part_2(self, inputs: Input) -> Output2:
+        boxes, pairs = inputs
+        dsu = DSU(len(boxes))
+        for pair in pairs:
+            dsu.unify(pair[0], pair[1])
+            if dsu.num_components == 1 and dsu.sz[dsu.find(0)] == len(boxes):
+                return boxes[pair[0]].x * boxes[pair[1]].x
         raise AssertionError
 
-    def part_2(self, pairs: Input) -> Output2:
-        return self.solve_2(pairs, 1000)
-
     def samples(self) -> None:
-        pairs = self.parse_input(TEST.splitlines())
-        assert self.solve_1(pairs, 10) == 40
-        assert self.solve_2(pairs, 10) == 25272
+        inputs = self.parse_input(TEST.splitlines())
+        assert self.solve_1(inputs, 10) == 40
+        assert self.part_2(inputs) == 25272
 
 
 solution = Solution(2025, 8)

src/main/python/aoc/geometry3d.py

@@ -1,8 +1,12 @@
 from __future__ import annotations
 
 import itertools
-from typing import Iterator
+from typing import TYPE_CHECKING
 from typing import NamedTuple
+from typing import Self
+
+if TYPE_CHECKING:
+    from collections.abc import Iterator
 
 from aoc.range import RangeInclusive
 
@@ -15,41 +19,46 @@ class Point3D(NamedTuple):
 
 class Position3D(Point3D):
     @classmethod
-    def copy(cls, position: Position3D) -> Position3D:
-        return Position3D.of(position.x, position.y, position.z)
+    def copy(cls, position: Self) -> Self:
+        return cls(position.x, position.y, position.z)
 
     @classmethod
-    def of(cls, x: int, y: int, z: int) -> Position3D:
-        return Position3D(x, y, z)
+    def of(cls, x: int, y: int, z: int) -> Self:
+        return cls(x, y, z)
 
     def translate(self, vector: Vector3D, amplitude: int = 1) -> Position3D:
-        return Position3D.of(
+        return Position3D(
             self.x + vector.x * amplitude,
             self.y + vector.y * amplitude,
             self.z + vector.z * amplitude,
         )
 
-    def manhattan_distance(self, other: Position3D) -> int:
+    def manhattan_distance(self, other: Self) -> int:
         return (
             abs(self.x - other.x)
             + abs(self.y - other.y)
             + abs(self.z - other.z)
         )
 
     def manhattan_distance_to_origin(self) -> int:
-        return self.manhattan_distance(Position3D.of(0, 0, 0))
+        return self.manhattan_distance(ORIGIN)
+
+    def squared_distance(self, other: Self) -> int:
+        dx, dy, dz = self.x - other.x, self.y - other.y, self.z - other.z
+        return dx * dx + dy * dy + dz * dz
+
+
+ORIGIN = Position3D(0, 0, 0)
 
 
 class Vector3D(Point3D):
     @classmethod
-    def of(cls, x: int, y: int, z: int) -> Vector3D:
-        return Vector3D(x, y, z)
+    def of(cls, x: int, y: int, z: int) -> Self:
+        return cls(x, y, z)
 
     @classmethod
-    def to_from(
-        cls, to: Position3D, from_: Position3D = Position3D.of(0, 0, 0)
-    ) -> Vector3D:
-        return Vector3D(to.x - from_.x, to.y - from_.y, to.z - from_.z)
+    def to_from(cls, to: Position3D, from_: Position3D = ORIGIN) -> Self:
+        return cls(to.x - from_.x, to.y - from_.y, to.z - from_.z)
 
 
 class Cuboid(NamedTuple):
@@ -61,10 +70,10 @@ class Cuboid(NamedTuple):
     z2: int
 
     @classmethod
-    def of(
+    def of(  # noqa:PLR0913
         cls, x1: int, x2: int, y1: int, y2: int, z1: int, z2: int
-    ) -> Cuboid:
-        return Cuboid(x1, x2, y1, y2, z1, z2)
+    ) -> Self:
+        return cls(x1, x2, y1, y2, z1, z2)
 
     def positions(self) -> Iterator[Position3D]:
         for (x, y), z in itertools.product(
@@ -75,17 +84,17 @@ def positions(self) -> Iterator[Position3D]:
         ):
             yield Position3D(x, y, z)
 
-    def overlap_x(self, other: Cuboid) -> bool:
+    def overlap_x(self, other: Self) -> bool:
         return RangeInclusive.between(self.x1, self.x2).is_overlapped_by(
             RangeInclusive.between(other.x1, other.x2)
         )
 
-    def overlap_y(self, other: Cuboid) -> bool:
+    def overlap_y(self, other: Self) -> bool:
         return RangeInclusive.between(self.y1, self.y2).is_overlapped_by(
             RangeInclusive.between(other.y1, other.y2)
         )
 
-    def overlap_z(self, other: Cuboid) -> bool:
+    def overlap_z(self, other: Self) -> bool:
         return RangeInclusive.between(self.z1, self.z2).is_overlapped_by(
             RangeInclusive.between(other.z1, other.z2)
         )