from functools import reduce from heapq import heappush from itertools import combinations Box = tuple[int, int, int] def distance_squared(left_box: Box, right_box: Box) -> int: return (right_box[0] - left_box[0])**2 + (right_box[1] - left_box[1])**2 + (right_box[2] - left_box[2])**2 def main(): junction_boxes: list[Box] with open("input_day8.txt") as f: junction_boxes = [tuple(map(int, line.split(","))) for line in f] n = 10 if len(junction_boxes) > 100: n = 1000 distances: list[tuple[int, Box, Box]] = [] for (left_box, right_box) in combinations(junction_boxes, 2): heappush(distances, (distance_squared(left_box, right_box), left_box, right_box)) circuits: dict[int, list[Box]] = {i: [b] for i, b in enumerate(junction_boxes)} connected_junction_boxes: dict[Box, int] = {b: i for i, b in enumerate(junction_boxes)} for i, (_, left_box, right_box) in enumerate(sorted(distances)): if i == n: print(reduce(lambda x, y: x * y, sorted(map(len, circuits.values()), reverse=True)[:3])) left_circuit_id = connected_junction_boxes.get(left_box) right_circuit_id = connected_junction_boxes.get(right_box) if left_circuit_id != right_circuit_id: for box in circuits[right_circuit_id]: connected_junction_boxes[box] = left_circuit_id circuits[left_circuit_id].extend(circuits[right_circuit_id]) del circuits[right_circuit_id] if len(circuits) == 1: print(left_box[0] * right_box[0]) return if __name__ == "__main__": main()