1
0

day8.py 2.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. import os
  2. import math
  3. with open(os.path.join(os.path.dirname(__file__), "example.txt")) as example:
  4. example_data = example.read().splitlines()
  5. with open(os.path.join(os.path.dirname(__file__), "input.txt")) as example:
  6. input_data = example.read().splitlines()
  7. def bench(part):
  8. import time
  9. def wrapper(*args, **kwargs):
  10. start = time.perf_counter()
  11. value = part(*args, **kwargs)
  12. print(f"\tevaluation time: {time.perf_counter() - start} s")
  13. return value
  14. return wrapper
  15. def cluster_by_distance(points, max_links=None):
  16. """
  17. points: list of list shape 3
  18. returns cluster as a dict, each cluster – list if index of points
  19. """
  20. n = len(points)
  21. # collect edges/distances
  22. edges = []
  23. for i in range(n):
  24. x1, y1, z1 = points[i]
  25. for j in range(i + 1, n):
  26. x2, y2, z2 = points[j]
  27. dx = x1 - x2
  28. dy = y1 - y2
  29. dz = z1 - z2
  30. dist = math.sqrt(dx * dx + dy * dy + dz * dz)
  31. edges.append((dist, i, j))
  32. edges.sort(key=lambda e: e[0])
  33. max_links = len(edges) if not max_links else max_links
  34. # Union–Find (DSU)
  35. parent = list(range(n))
  36. rank = [0] * n
  37. def find(x):
  38. while parent[x] != x:
  39. parent[x] = parent[parent[x]]
  40. x = parent[x]
  41. return x
  42. def union(a, b):
  43. ra, rb = find(a), find(b)
  44. if ra == rb:
  45. return False
  46. if rank[ra] < rank[rb]:
  47. parent[ra] = rb
  48. elif rank[ra] > rank[rb]:
  49. parent[rb] = ra
  50. else:
  51. parent[rb] = ra
  52. rank[ra] += 1
  53. return True
  54. # just walk through the edges
  55. links_done = 0
  56. for dist, i, j in edges[:max_links]:
  57. if union(i, j):
  58. last_pair = i, j
  59. links_done += 1
  60. # collect clusters
  61. clusters = {}
  62. for i in range(n):
  63. r = find(i)
  64. clusters.setdefault(r, []).append(i)
  65. return clusters, links_done, last_pair
  66. def part1(data, n):
  67. from functools import reduce
  68. from operator import mul
  69. points = [list(map(int, point.split(","))) for point in data]
  70. clusters, _, _ = cluster_by_distance(points, n)
  71. result = reduce(mul, [len(c[1]) for c in sorted(clusters.items(), key=lambda x: len(x[1]), reverse=True)[:3]])
  72. print(f"Part1: {result}")
  73. @bench
  74. def part2(data):
  75. points = [list(map(int, point.split(","))) for point in data]
  76. _, _, last_pair = cluster_by_distance(points)
  77. result = points[last_pair[0]][0] * points[last_pair[1]][0]
  78. print(f"Part2: {result}")
  79. part1(example_data, 10)
  80. part1(input_data, 1000)
  81. part2(example_data)
  82. part2(input_data)