kruskal

kruskal.py

E = {
    "A": [("B", 4), ("H", 8)],
    "B": [("A", 4), ("C", 8), ("H", 11)],
    "C": [("B", 8), ("D", 7), ("F", 4), ("I", 2)],
    "D": [("C", 7), ("E", 9), ("F", 14)],
    "E": [("D", 9), ("F", 10)],
    "F": [("C", 4), ("D", 14), ("E", 10), ("G", 2)],
    "G": [("F", 2), ("H", 1), ("I", 6)],
    "H": [("A", 8), ("B", 11), ("G", 1), ("I", 7)],
    "I": [("C", 2), ("G", 6), ("H", 7)],
}

from union_find import WeightedQuickUnion
import heapq

class Kruskal:
    def __init__(self, graph):
        self.graph = graph

    def run(self):
        edges_seen = set()
        edge_heap = []
        for u in self.graph:
            for v, w in self.graph[u]:
                # Avoid duplicates in undirected graph
                if (v, u) not in edges_seen and (u, v) not in edges_seen:
                    heapq.heappush(edge_heap, (w, u, v))
                    edges_seen.add((u, v))
        uf = WeightedQuickUnion(self.graph.keys())
        MST = []
        n_nodes = len(self.graph)
        while edge_heap and len(MST) < n_nodes - 1:
            w, u, v = heapq.heappop(edge_heap)
            if not uf.connected(u, v):
                uf.union(u, v)
                MST.append((u, v, w))
        return MST
    
if __name__ == "__main__":
    kruskal = Kruskal(E)
    MST = kruskal.run()
    print("Edges in the Minimum Spanning Tree:")
    for u, v, w in MST:
        print(f"{u} -- {v} (weight {w})")
    total_weight = sum(w for _, _, w in MST)
    print(f"Total weight of MST: {total_weight}")

    # visualizating the MST using networkx
    import networkx as nx
    import matplotlib.pyplot as plt
    G = nx.Graph()
    for u in E:
        for v, w in E[u]:
            G.add_edge(u, v, weight=w)
    pos = nx.spring_layout(G)
    nx.draw(G, pos, with_labels=True, node_color='lightblue', node_size=500)
    edge_labels = nx.get_edge_attributes(G, 'weight')
    nx.draw_networkx_edge_labels(G, pos, edge_labels=edge_labels)
    # highlight the MST edges
    mst_edges = [(u, v) for u, v, w in MST]
    nx.draw_networkx_edges(G, pos, edgelist=mst_edges, edge_color='r', width=2)
    plt.title("Graph with Minimum Spanning Tree Highlighted")
    plt.show()