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