Python
Both of my solutions make use of a Heap and a Disjoint-Set-Union / Union-Find data structure
Part 1: Use a heap to keep 1000 shortest connections, then union them all in DSU and get the 3 largest connected components.
Part 2: Use a heap to keep all connections, then union the shortest ones until all components are connected.
click to view code
import heapq as hq
from itertools import combinations
from collections import Counter
# Disjoint Set Union (or Union-Find) data structure
# with path compression and union by rank
class DSU:
def __init__(self, size: int) -> None:
self.parent = list(range(size)) # ith value is the parent node to node i
self.rank = [1] * size # max depth of subtree rooted here (used for union by rank)
def find(self, x: int) -> int:
# if the node is not its own parent, we need to recurse on parent
if x != self.parent[x]:
# path compression
self.parent[x] = self.find(self.parent[x])
return self.parent[x]
def isConnected(self, x: int, y: int) -> bool:
return self.find(x) == self.find(y)
# returns a boolean whether or not union was needed
def union(self, x: int, y: int) -> bool:
rootX = self.find(x)
rootY = self.find(y)
if rootX == rootY:
# no union needed
return False
if self.rank[rootX] > self.rank[rootY]:
# rootX has deeper subtree, therefore set it as parent to rootY (and its subtree)
self.parent[rootY] = rootX
elif self.rank[rootX] < self.rank[rootY]:
# rootY has deeper subtree, therefore set it as parent to rootX (and its subtree)
self.parent[rootX] = rootY
else:
# both subtrees are of equal depth, therefore choose either one of them as the parent of the other
# here we chose rootX as the parent of rootY, therefore rootX depth increases by 1
self.parent[rootY] = rootX
self.rank[rootX] += 1
# union complete
return True
# parse puzzle input into junction coordinates (x, y, z)
def parse_junctions(data: str):
return [tuple(map(int, line.split(','))) for line in data.splitlines()]
# get distance between two points in 3D space
# skipping the sqrt because we only care about relative distances
def get_dist(p1, p2):
x1, y1, z1 = p1
x2, y2, z2 = p2
return (x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2
def part1(data: str, max_connections: int = 1000):
junctions = parse_junctions(data)
n = len(junctions)
# <max_connections> shortest connections
closest_connections = []
# iterate over all pairs of junctions
# combinations is just some syntactic sugar to avoid a nested for loop (still O(n^2))
for i, j in combinations(range(n), 2):
dist = get_dist(junctions[i], junctions[j])
# keep <max_connections> closest connections
# have to use -dist to simulate max heap
if len(closest_connections) < max_connections:
hq.heappush(closest_connections, (-dist, i, j))
else:
hq.heappushpop(closest_connections, (-dist, i, j))
# union all the closest connections
dsu = DSU(n)
for _, i, j in closest_connections:
dsu.union(i, j)
# count all circuit sizes
circuit_sizes = Counter()
for i in range(n):
circuit_sizes[dsu.find(i)] += 1
# multiply the sizes of the 3 largest circuits
ans = 1
for _, f in circuit_sizes.most_common(3):
ans *= f
return ans
def part2(data: str):
junctions = parse_junctions(data)
n = len(junctions)
# all n^2 junction connections
all_connections = []
# combinations is just some syntactic sugar to avoid a nested for loop (still O(n^2))
for i, j in combinations(range(n), 2):
dist = get_dist(junctions[i], junctions[j])
all_connections.append((dist, i, j))
# create min heap of all connections
# Heapify later to benefit from the heapify algorithm (O(n) instead of O(n log n))
hq.heapify(all_connections)
# connect junctions until all are connected
dsu = DSU(n)
unions = 0
while all_connections:
_, i, j = hq.heappop(all_connections)
unions += dsu.union(i, j)
# if we have n-1 unions, that means all n junctions are connected
if unions == n-1:
return junctions[i][0] * junctions[j][0]
assert False, "It's not possible that all junctions never connect"
sample = """162,817,812
57,618,57
906,360,560
592,479,940
352,342,300
466,668,158
542,29,236
431,825,988
739,650,466
52,470,668
216,146,977
819,987,18
117,168,530
805,96,715
346,949,466
970,615,88
941,993,340
862,61,35
984,92,344
425,690,689"""
assert part1(sample, max_connections=10) == 40
assert part2(sample) == 25272
