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