1"""
2Minimum Spanning Tree - Kruskal's Algorithm
3
4Kruskal's algorithm builds the MST by sorting all edges by weight and adding
5them one by one, skipping edges that would create a cycle. Uses Union-Find
6(Disjoint Set) to efficiently detect cycles.
7
8Key Insight:
9 Process edges globally from smallest to largest. An edge is safe to add
10 if it connects two different components (detected via Union-Find).
11
12Time Complexity: O(E log E) - dominated by edge sorting
13Space Complexity: O(V) - Union-Find data structure
14
15Advantages:
16 - Works well with sparse graphs
17 - Natural for edge list representation
18 - Easy to parallelize (sorting step)
19 - Finds MST forest for disconnected graphs
20
21Graph Representation:
22 Uses adjacency list format for WEIGHTED UNDIRECTED graphs.
23 Format: {node: {neighbor: weight, ...}}
24"""
25
26from typing import Dict, Set, List, Any, Tuple
27
28
29class UnionFind:
30 """Union-Find with path compression and union by rank."""
31
32 def __init__(self, elements: Set[Any]):
33 self.parent = {x: x for x in elements}
34 self.rank = {x: 0 for x in elements}
35
36 def find(self, x: Any) -> Any:
37 """Find root with path compression."""
38 if self.parent[x] != x:
39 self.parent[x] = self.find(self.parent[x])
40 return self.parent[x]
41
42 def union(self, x: Any, y: Any) -> bool:
43 """Union by rank. Returns True if merged, False if already connected."""
44 root_x, root_y = self.find(x), self.find(y)
45
46 if root_x == root_y:
47 return False
48
49 if self.rank[root_x] < self.rank[root_y]:
50 root_x, root_y = root_y, root_x
51
52 self.parent[root_y] = root_x
53 if self.rank[root_x] == self.rank[root_y]:
54 self.rank[root_x] += 1
55
56 return True
57
58
59def mst(graph: Dict[Any, Dict[Any, float]]) -> List[Tuple[Any, Any, float]]:
60 """
61 Find the Minimum Spanning Tree using Kruskal's algorithm.
62
63 Args:
64 graph: Weighted adjacency list {node: {neighbor: weight, ...}}
65
66 Returns:
67 List of edges in MST as (from, to, weight) tuples
68
69 Raises:
70 ValueError: If graph is empty or disconnected
71 """
72 if not graph:
73 raise ValueError("Graph is empty")
74
75 # Extract all edges (avoid duplicates for undirected graph)
76 edges = []
77 seen = set()
78 for node in graph:
79 for neighbor, weight in graph[node].items():
80 edge_key = (min(node, neighbor), max(node, neighbor))
81 if edge_key not in seen:
82 edges.append((node, neighbor, weight))
83 seen.add(edge_key)
84
85 # Sort edges by weight
86 edges.sort(key=lambda e: e[2])
87
88 # Initialize Union-Find
89 uf = UnionFind(set(graph.keys()))
90
91 mst_edges = []
92 for frm, to, weight in edges:
93 if uf.union(frm, to): # Returns True if they were in different sets
94 mst_edges.append((frm, to, weight))
95
96 if len(mst_edges) == len(graph) - 1:
97 break # MST complete
98
99 if len(mst_edges) != len(graph) - 1:
100 raise ValueError("Graph is disconnected - MST does not exist")
101
102 return mst_edges
103
104
105def mst_weight(graph: Dict[Any, Dict[Any, float]]) -> float:
106 """
107 Calculate the total weight of the Minimum Spanning Tree.
108
109 Args:
110 graph: Weighted adjacency list
111
112 Returns:
113 Total weight of the MST
114 """
115 edges = mst(graph)
116 return sum(weight for _, _, weight in edges)
117
118
119def mst_as_graph(graph: Dict[Any, Dict[Any, float]]) -> Dict[Any, Dict[Any, float]]:
120 """
121 Return the MST as an adjacency list (undirected).
122
123 Args:
124 graph: Weighted adjacency list
125
126 Returns:
127 MST as weighted adjacency list
128 """
129 edges = mst(graph)
130 result: Dict[Any, Dict[Any, float]] = {node: {} for node in graph}
131
132 for frm, to, weight in edges:
133 result[frm][to] = weight
134 result[to][frm] = weight
135
136 return result
137
138
139# Example usage and demonstration
140if __name__ == "__main__":
141 # Example weighted undirected graph
142 #
143 # A ---4--- B
144 # | \ |
145 # 2 3 1
146 # | \ |
147 # C ---5--- D
148 #
149 example_graph = {
150 "A": {"B": 4, "C": 2, "D": 3},
151 "B": {"A": 4, "D": 1},
152 "C": {"A": 2, "D": 5},
153 "D": {"A": 3, "B": 1, "C": 5},
154 }
155
156 print("MST Edges (Kruskal):", mst(example_graph))
157 print("MST Total Weight:", mst_weight(example_graph))
158 print("MST as Graph:", mst_as_graph(example_graph))
159
160 # Disconnected graph (should raise error)
161 disconnected = {
162 "A": {"B": 1},
163 "B": {"A": 1},
164 "C": {"D": 2},
165 "D": {"C": 2},
166 }
167
168 try:
169 mst(disconnected)
170 except ValueError as e:
171 print(f"Disconnected graph error: {e}")
172