Day8.py 1.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. from functools import reduce
  2. from heapq import heappush
  3. from itertools import combinations
  4. Box = tuple[int, int, int]
  5. def distance_squared(left_box: Box, right_box: Box) -> int:
  6. return (right_box[0] - left_box[0])**2 + (right_box[1] - left_box[1])**2 + (right_box[2] - left_box[2])**2
  7. def main():
  8. junction_boxes: list[Box]
  9. with open("input_day8.txt") as f:
  10. junction_boxes = [tuple(map(int, line.split(","))) for line in f]
  11. n = 10
  12. if len(junction_boxes) > 100:
  13. n = 1000
  14. distances: list[tuple[int, Box, Box]] = []
  15. for (left_box, right_box) in combinations(junction_boxes, 2):
  16. heappush(distances, (distance_squared(left_box, right_box), left_box, right_box))
  17. circuits: dict[int, list[Box]] = {i: [b] for i, b in enumerate(junction_boxes)}
  18. connected_junction_boxes: dict[Box, int] = {b: i for i, b in enumerate(junction_boxes)}
  19. for i, (_, left_box, right_box) in enumerate(sorted(distances)):
  20. if i == n:
  21. print(reduce(lambda x, y: x * y, sorted(map(len, circuits.values()), reverse=True)[:3]))
  22. left_circuit_id = connected_junction_boxes.get(left_box)
  23. right_circuit_id = connected_junction_boxes.get(right_box)
  24. if left_circuit_id != right_circuit_id:
  25. for box in circuits[right_circuit_id]:
  26. connected_junction_boxes[box] = left_circuit_id
  27. circuits[left_circuit_id].extend(circuits[right_circuit_id])
  28. del circuits[right_circuit_id]
  29. if len(circuits) == 1:
  30. print(left_box[0] * right_box[0])
  31. return
  32. if __name__ == "__main__":
  33. main()