101 lines
2.9 KiB
Python
101 lines
2.9 KiB
Python
"""Kruskal's minimum spanning tree on an undirected weighted graph.
|
|
|
|
Public API:
|
|
kruskal(edges, nodes) -> (total_weight, list_of_edges)
|
|
edges: iterable of (u, v, w) (undirected). nodes: iterable of node labels.
|
|
Returns (weight, edges) of an MST; raises ValueError if graph is disconnected
|
|
(no spanning tree exists). Ties broken by input order (stable).
|
|
|
|
Also exposes `find(parent, x)` / `union(parent, rank, x, y)` helpers.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
|
|
def find(parent, x):
|
|
# iterative path compression
|
|
root = x
|
|
while parent[root] != root:
|
|
root = parent[root]
|
|
while parent[x] != root:
|
|
parent[x], x = root, parent[x]
|
|
return root
|
|
|
|
|
|
def union(parent, rank, x, y):
|
|
rx, ry = find(parent, x), find(parent, y)
|
|
if rx == ry:
|
|
return False
|
|
if rank[rx] < rank[ry]:
|
|
rx, ry = ry, rx
|
|
parent[ry] = rx
|
|
if rank[rx] == rank[ry]:
|
|
rank[rx] += 1
|
|
return True
|
|
|
|
|
|
def kruskal(edges, nodes):
|
|
nodes = list(nodes)
|
|
parent = {n: n for n in nodes}
|
|
rank = {n: 0 for n in nodes}
|
|
# stable sort by weight (Python sort is stable)
|
|
ordered = sorted(edges, key=lambda e: e[2])
|
|
chosen = []
|
|
total = 0
|
|
for u, v, w in ordered:
|
|
if u not in parent or v not in parent:
|
|
raise ValueError("unknown node")
|
|
if w < 0:
|
|
raise ValueError("negative weight")
|
|
if union(parent, rank, u, v):
|
|
chosen.append((u, v, w))
|
|
total += w
|
|
if len(chosen) == len(nodes) - 1:
|
|
break
|
|
if len(nodes) > 1 and len(chosen) != len(nodes) - 1:
|
|
raise ValueError("graph disconnected")
|
|
return total, chosen
|
|
|
|
|
|
def _brute_mst(edges, nodes):
|
|
"""Reference: enumerate all subsets of edges of size |V|-1, pick min weight
|
|
spanning tree (connected, acyclic). Only feasible for small graphs."""
|
|
from itertools import combinations
|
|
|
|
nodes = list(nodes)
|
|
n = len(nodes)
|
|
if n <= 1:
|
|
return 0, []
|
|
need = n - 1
|
|
best = None
|
|
for combo in combinations(edges, need):
|
|
# build adjacency
|
|
adj = {x: [] for x in nodes}
|
|
for u, v, w in combo:
|
|
adj[u].append(v)
|
|
adj[v].append(u)
|
|
# connectivity
|
|
seen = set()
|
|
stack = [nodes[0]]
|
|
while stack:
|
|
x = stack.pop()
|
|
if x in seen:
|
|
continue
|
|
seen.add(x)
|
|
for y in adj[x]:
|
|
if y not in seen:
|
|
stack.append(y)
|
|
if len(seen) != n:
|
|
continue
|
|
# acyclicity: edges == nodes-1 and connected => tree
|
|
total = sum(w for _, _, w in combo)
|
|
if best is None or total < best[0]:
|
|
best = (total, list(combo))
|
|
if best is None:
|
|
raise ValueError("graph disconnected")
|
|
return best
|
|
|
|
|
|
if __name__ == "__main__":
|
|
e = [("a", "b", 1), ("b", "c", 2), ("a", "c", 4)]
|
|
print(kruskal(e, ["a", "b", "c"]))
|