beam: Kruskal MST with union-find, brute-force-verified

wizard-work/beam/kruskal.py

@@ -0,0 +1,100 @@
+"""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"]))

wizard-work/beam/test_kruskal.py

@@ -0,0 +1,72 @@
+import random
+from kruskal import kruskal, _brute_mst
+
+
+def test_basic():
+    e = [("a", "b", 1), ("b", "c", 2), ("a", "c", 4)]
+    w, chosen = kruskal(e, ["a", "b", "c"])
+    assert w == 3
+    assert len(chosen) == 2
+
+
+def test_single_node():
+    w, chosen = kruskal([], ["a"])
+    assert w == 0 and chosen == []
+
+
+def test_disconnected():
+    try:
+        kruskal([("a", "b", 1)], ["a", "b", "c"])
+        raise AssertionError("expected")
+    except ValueError:
+        pass
+
+
+def test_negative():
+    try:
+        kruskal([("a", "b", -1)], ["a", "b"])
+        raise AssertionError("expected")
+    except ValueError:
+        pass
+
+
+def test_unknown_node():
+    try:
+        kruskal([("a", "z", 1)], ["a", "b"])
+        raise AssertionError("expected")
+    except ValueError:
+        pass
+
+
+def test_random_vs_brute():
+    rng = random.Random(42)
+    for _ in range(300):
+        n = rng.randint(1, 6)
+        nodes = list(range(n))
+        # random edge set
+        possible = [(u, v) for u in range(n) for v in range(u + 1, n)]
+        m = rng.randint(0, len(possible))
+        edges = []
+        for u, v in rng.sample(possible, m):
+            edges.append((u, v, rng.randint(0, 9)))
+        # maybe duplicate weights fine
+        try:
+            w1, _ = kruskal(edges, nodes)
+            ok1 = True
+        except ValueError:
+            ok1 = False
+        try:
+            w2, _ = _brute_mst(edges, nodes)
+            ok2 = True
+        except ValueError:
+            ok2 = False
+        assert ok1 == ok2, (edges, nodes, ok1, ok2)
+        if ok1:
+            assert w1 == w2, (edges, nodes, w1, w2)
+
+
+if __name__ == "__main__":
+    for n in [x for x in globals() if x.startswith("test_")]:
+        globals()[n]()
+        print("PASS", n)
+    print("ALL PASS")