kruskal.py

# UnionFind is defined in https://gist.github.com/siavashk/fef678f9f2ef17fdc6cff917922bf4ba
# Augment the union function to return True for new connections
def kruskal(edges: list[tuple[int, int, int]]):
    """ 
    Edges: list of edge in (weight, start, end]) format

    """
    edges.sort()
    uf = UnionFind()
    for _, x, y in edges:
        uf.add(x)
        uf.add(y)
    
    total_cost = 0
    mst_edges = []
    for w, x, y in edges:
        if uf.union(x, y):
            total_cost += w
            mst_edges.append([w, x, y])
        if len(mst_edges) == len(uf.parent) - 1:
            break
    
    return total_cost, mst_edges