Skip to main content

sim_lib_discrete_graph/
mst.rs

1//! Minimum spanning tree via Prim's and Kruskal's algorithms.
2//!
3//! Both require an undirected graph, ignore self-loops, and break ties
4//! deterministically by `(weight, min endpoint, max endpoint, edge id)`. Each
5//! returns a [`SpanningTree`] whose edge ids are sorted ascending, so the two
6//! algorithms produce comparable witnesses.
7
8use crate::certificate::SpanningTree;
9use crate::error::GraphError;
10use crate::graph::Graph;
11use crate::unionfind::UnionFind;
12use core::cmp::Reverse;
13use core::ops::Add;
14use std::collections::BinaryHeap;
15
16/// Prim heap key: `(weight, min endpoint, max endpoint, edge id, to-node)`.
17type PrimHeap<W> = BinaryHeap<Reverse<(W, usize, usize, usize, usize)>>;
18
19fn require_undirected<N, W>(graph: &Graph<N, W>) -> Result<(), GraphError> {
20    graph.validate()?;
21    if graph.is_directed() {
22        return Err(GraphError::WrongGraphKind(
23            "MST requires an undirected graph".to_string(),
24        ));
25    }
26    Ok(())
27}
28
29/// Kruskal's algorithm: sort edges, union-find to reject cycles.
30pub fn kruskals_mst<N, W>(graph: &Graph<N, W>) -> Result<SpanningTree<W>, GraphError>
31where
32    W: Ord + Clone + Default + Add<Output = W>,
33{
34    require_undirected(graph)?;
35    let n = graph.node_count();
36    if n == 0 {
37        return Ok(SpanningTree {
38            edges: Vec::new(),
39            total_weight: W::default(),
40        });
41    }
42    let mut cand: Vec<_> = graph.edges.iter().filter(|e| !e.is_self_loop()).collect();
43    cand.sort_by(|a, b| {
44        a.weight
45            .cmp(&b.weight)
46            .then_with(|| a.source.min(a.target).cmp(&b.source.min(b.target)))
47            .then_with(|| a.source.max(a.target).cmp(&b.source.max(b.target)))
48            .then_with(|| a.id.cmp(&b.id))
49    });
50    let mut uf = UnionFind::new(n);
51    let mut chosen = Vec::new();
52    let mut total = W::default();
53    for e in cand {
54        if uf.union(e.source, e.target) {
55            chosen.push(e.id);
56            total = total + e.weight.clone();
57            if chosen.len() == n - 1 {
58                break;
59            }
60        }
61    }
62    if chosen.len() != n - 1 {
63        return Err(GraphError::Disconnected);
64    }
65    chosen.sort_unstable();
66    Ok(SpanningTree {
67        edges: chosen,
68        total_weight: total,
69    })
70}
71
72/// Prim's algorithm: grow a tree from node 0 using a min-heap keyed on the same
73/// deterministic tie-break as Kruskal.
74pub fn prims_mst<N, W>(graph: &Graph<N, W>) -> Result<SpanningTree<W>, GraphError>
75where
76    W: Ord + Clone + Default + Add<Output = W>,
77{
78    require_undirected(graph)?;
79    let n = graph.node_count();
80    if n == 0 {
81        return Ok(SpanningTree {
82            edges: Vec::new(),
83            total_weight: W::default(),
84        });
85    }
86    // adj[x] = (weight, other endpoint, edge id), excluding self-loops.
87    let mut adj: Vec<Vec<(W, usize, usize)>> = vec![Vec::new(); n];
88    for e in &graph.edges {
89        if e.is_self_loop() {
90            continue;
91        }
92        adj[e.source].push((e.weight.clone(), e.target, e.id));
93        adj[e.target].push((e.weight.clone(), e.source, e.id));
94    }
95
96    let mut in_tree = vec![false; n];
97    let mut heap: PrimHeap<W> = BinaryHeap::new();
98    let push_incident = |heap: &mut PrimHeap<W>, in_tree: &[bool], x: usize| {
99        for (w, other, id) in &adj[x] {
100            if !in_tree[*other] {
101                heap.push(Reverse((
102                    w.clone(),
103                    x.min(*other),
104                    x.max(*other),
105                    *id,
106                    *other,
107                )));
108            }
109        }
110    };
111
112    in_tree[0] = true;
113    push_incident(&mut heap, &in_tree, 0);
114    let mut chosen = Vec::new();
115    let mut total = W::default();
116    let mut count = 1;
117    while count < n {
118        let Some(Reverse((w, _, _, id, to))) = heap.pop() else {
119            break;
120        };
121        if in_tree[to] {
122            continue;
123        }
124        in_tree[to] = true;
125        chosen.push(id);
126        total = total + w;
127        count += 1;
128        push_incident(&mut heap, &in_tree, to);
129    }
130    if count != n {
131        return Err(GraphError::Disconnected);
132    }
133    chosen.sort_unstable();
134    Ok(SpanningTree {
135        edges: chosen,
136        total_weight: total,
137    })
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143    use crate::certificate::verify_mst;
144    use crate::edge::Directedness;
145
146    fn triangle() -> Graph<u8, u64> {
147        let mut g = Graph::with_nodes(vec![0, 1, 2], Directedness::Undirected);
148        g.add_edge(0, 1, 1).unwrap(); // id 0
149        g.add_edge(1, 2, 2).unwrap(); // id 1
150        g.add_edge(0, 2, 3).unwrap(); // id 2
151        g
152    }
153
154    #[test]
155    fn triangle_mst_weight_and_edges() {
156        let t = kruskals_mst(&triangle()).unwrap();
157        assert_eq!(t.total_weight, 3);
158        assert_eq!(t.edges, vec![0, 1]);
159    }
160
161    #[test]
162    fn prim_equals_kruskal_weight() {
163        let g = triangle();
164        assert_eq!(
165            prims_mst(&g).unwrap().total_weight,
166            kruskals_mst(&g).unwrap().total_weight
167        );
168    }
169
170    #[test]
171    fn equal_weight_ties_are_deterministic() {
172        // All weight 1; the tie-break selects the lowest (min,max,id) edges.
173        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2, 3], Directedness::Undirected);
174        g.add_edge(0, 1, 1).unwrap(); // id 0
175        g.add_edge(1, 2, 1).unwrap(); // id 1
176        g.add_edge(2, 3, 1).unwrap(); // id 2
177        g.add_edge(0, 3, 1).unwrap(); // id 3
178        g.add_edge(0, 2, 1).unwrap(); // id 4
179        assert_eq!(kruskals_mst(&g).unwrap().edges, vec![0, 3, 4]);
180    }
181
182    #[test]
183    fn disconnected_graph_fails() {
184        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2, 3], Directedness::Undirected);
185        g.add_edge(0, 1, 1).unwrap();
186        assert_eq!(kruskals_mst(&g).unwrap_err(), GraphError::Disconnected);
187    }
188
189    #[test]
190    fn directed_graph_is_wrong_kind() {
191        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1], Directedness::Directed);
192        g.add_edge(0, 1, 1).unwrap();
193        assert!(matches!(
194            kruskals_mst(&g),
195            Err(GraphError::WrongGraphKind(_))
196        ));
197    }
198
199    #[test]
200    fn valid_certificate_verifies_tampered_rejected() {
201        let g = triangle();
202        let cert = kruskals_mst(&g).unwrap().certificate();
203        assert!(verify_mst(&g, &cert).is_ok());
204
205        // Tamper with the weight.
206        let mut bad = cert.clone();
207        bad.total_weight_repr = "99".to_string();
208        assert!(verify_mst(&g, &bad).is_err());
209
210        // A suboptimal spanning tree {edge1, edge2} (weight 5) is rejected by
211        // the cycle property: edge0 (weight 1) is cheaper than the path max 3.
212        let suboptimal = crate::certificate::MstCertificate {
213            edge_ids: vec![1, 2],
214            total_weight_repr: "5".to_string(),
215        };
216        assert!(verify_mst(&g, &suboptimal).is_err());
217    }
218}