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 std::collections::BinaryHeap;
14
15/// Prim heap key: `(weight, min endpoint, max endpoint, edge id, to-node)`.
16type PrimHeap<W> = BinaryHeap<Reverse<(W, usize, usize, usize, usize)>>;
17
18/// Integer weights that can report overflow while summing MST totals.
19pub trait MstWeight: Ord + Clone + Default {
20    /// Returns `self + rhs`, or `None` when the sum overflows.
21    fn checked_add(&self, rhs: &Self) -> Option<Self>;
22}
23
24macro_rules! impl_mst_weight {
25    ($($ty:ty),* $(,)?) => {
26        $(
27            impl MstWeight for $ty {
28                fn checked_add(&self, rhs: &Self) -> Option<Self> {
29                    <$ty>::checked_add(*self, *rhs)
30                }
31            }
32        )*
33    };
34}
35
36impl_mst_weight!(
37    i8, i16, i32, i64, i128, isize, u8, u16, u32, u64, u128, usize,
38);
39
40pub(crate) fn checked_mst_add<W: MstWeight>(left: &W, right: &W) -> Result<W, GraphError> {
41    left.checked_add(right)
42        .ok_or_else(|| GraphError::WeightOverflow("minimum spanning tree total".to_string()))
43}
44
45fn require_undirected<N, W>(graph: &Graph<N, W>) -> Result<(), GraphError> {
46    graph.validate()?;
47    if graph.is_directed() {
48        return Err(GraphError::WrongGraphKind(
49            "MST requires an undirected graph".to_string(),
50        ));
51    }
52    Ok(())
53}
54
55/// Kruskal's algorithm: sort edges, union-find to reject cycles.
56pub fn kruskals_mst<N, W>(graph: &Graph<N, W>) -> Result<SpanningTree<W>, GraphError>
57where
58    W: MstWeight,
59{
60    require_undirected(graph)?;
61    let n = graph.node_count();
62    if n == 0 {
63        return Ok(SpanningTree {
64            edges: Vec::new(),
65            total_weight: W::default(),
66        });
67    }
68    let mut cand: Vec<_> = graph.edges.iter().filter(|e| !e.is_self_loop()).collect();
69    cand.sort_by(|a, b| {
70        a.weight
71            .cmp(&b.weight)
72            .then_with(|| a.source.min(a.target).cmp(&b.source.min(b.target)))
73            .then_with(|| a.source.max(a.target).cmp(&b.source.max(b.target)))
74            .then_with(|| a.id.cmp(&b.id))
75    });
76    let mut uf = UnionFind::new(n);
77    let mut chosen = Vec::new();
78    let mut total = W::default();
79    for e in cand {
80        if uf.union(e.source, e.target) {
81            chosen.push(e.id);
82            total = checked_mst_add(&total, &e.weight)?;
83            if chosen.len() == n - 1 {
84                break;
85            }
86        }
87    }
88    if chosen.len() != n - 1 {
89        return Err(GraphError::Disconnected);
90    }
91    chosen.sort_unstable();
92    Ok(SpanningTree {
93        edges: chosen,
94        total_weight: total,
95    })
96}
97
98/// Prim's algorithm: grow a tree from node 0 using a min-heap keyed on the same
99/// deterministic tie-break as Kruskal.
100pub fn prims_mst<N, W>(graph: &Graph<N, W>) -> Result<SpanningTree<W>, GraphError>
101where
102    W: MstWeight,
103{
104    require_undirected(graph)?;
105    let n = graph.node_count();
106    if n == 0 {
107        return Ok(SpanningTree {
108            edges: Vec::new(),
109            total_weight: W::default(),
110        });
111    }
112    // adj[x] = (weight, other endpoint, edge id), excluding self-loops.
113    let mut adj: Vec<Vec<(W, usize, usize)>> = vec![Vec::new(); n];
114    for e in &graph.edges {
115        if e.is_self_loop() {
116            continue;
117        }
118        adj[e.source].push((e.weight.clone(), e.target, e.id));
119        adj[e.target].push((e.weight.clone(), e.source, e.id));
120    }
121
122    let mut in_tree = vec![false; n];
123    let mut heap: PrimHeap<W> = BinaryHeap::new();
124    let push_incident = |heap: &mut PrimHeap<W>, in_tree: &[bool], x: usize| {
125        for (w, other, id) in &adj[x] {
126            if !in_tree[*other] {
127                heap.push(Reverse((
128                    w.clone(),
129                    x.min(*other),
130                    x.max(*other),
131                    *id,
132                    *other,
133                )));
134            }
135        }
136    };
137
138    in_tree[0] = true;
139    push_incident(&mut heap, &in_tree, 0);
140    let mut chosen = Vec::new();
141    let mut total = W::default();
142    let mut count = 1;
143    while count < n {
144        let Some(Reverse((w, _, _, id, to))) = heap.pop() else {
145            break;
146        };
147        if in_tree[to] {
148            continue;
149        }
150        in_tree[to] = true;
151        chosen.push(id);
152        total = checked_mst_add(&total, &w)?;
153        count += 1;
154        push_incident(&mut heap, &in_tree, to);
155    }
156    if count != n {
157        return Err(GraphError::Disconnected);
158    }
159    chosen.sort_unstable();
160    Ok(SpanningTree {
161        edges: chosen,
162        total_weight: total,
163    })
164}
165
166#[cfg(test)]
167mod tests {
168    use super::*;
169    use crate::certificate::verify_mst;
170    use crate::edge::Directedness;
171
172    fn triangle() -> Graph<u8, u64> {
173        let mut g = Graph::with_nodes(vec![0, 1, 2], Directedness::Undirected);
174        g.add_edge(0, 1, 1).unwrap(); // id 0
175        g.add_edge(1, 2, 2).unwrap(); // id 1
176        g.add_edge(0, 2, 3).unwrap(); // id 2
177        g
178    }
179
180    #[test]
181    fn triangle_mst_weight_and_edges() {
182        let t = kruskals_mst(&triangle()).unwrap();
183        assert_eq!(t.total_weight, 3);
184        assert_eq!(t.edges, vec![0, 1]);
185    }
186
187    #[test]
188    fn prim_equals_kruskal_weight() {
189        let g = triangle();
190        assert_eq!(
191            prims_mst(&g).unwrap().total_weight,
192            kruskals_mst(&g).unwrap().total_weight
193        );
194    }
195
196    #[test]
197    fn equal_weight_ties_are_deterministic() {
198        // All weight 1; the tie-break selects the lowest (min,max,id) edges.
199        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2, 3], Directedness::Undirected);
200        g.add_edge(0, 1, 1).unwrap(); // id 0
201        g.add_edge(1, 2, 1).unwrap(); // id 1
202        g.add_edge(2, 3, 1).unwrap(); // id 2
203        g.add_edge(0, 3, 1).unwrap(); // id 3
204        g.add_edge(0, 2, 1).unwrap(); // id 4
205        assert_eq!(kruskals_mst(&g).unwrap().edges, vec![0, 3, 4]);
206    }
207
208    #[test]
209    fn disconnected_graph_fails() {
210        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2, 3], Directedness::Undirected);
211        g.add_edge(0, 1, 1).unwrap();
212        assert_eq!(kruskals_mst(&g).unwrap_err(), GraphError::Disconnected);
213    }
214
215    #[test]
216    fn directed_graph_is_wrong_kind() {
217        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1], Directedness::Directed);
218        g.add_edge(0, 1, 1).unwrap();
219        assert!(matches!(
220            kruskals_mst(&g),
221            Err(GraphError::WrongGraphKind(_))
222        ));
223    }
224
225    #[test]
226    fn mst_total_weight_overflow_is_rejected() {
227        let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2], Directedness::Undirected);
228        g.add_edge(0, 1, u64::MAX).unwrap();
229        g.add_edge(1, 2, 1).unwrap();
230
231        assert!(matches!(
232            kruskals_mst(&g),
233            Err(GraphError::WeightOverflow(_))
234        ));
235        assert!(matches!(prims_mst(&g), Err(GraphError::WeightOverflow(_))));
236    }
237
238    #[test]
239    fn valid_certificate_verifies_tampered_rejected() {
240        let g = triangle();
241        let cert = kruskals_mst(&g).unwrap().certificate();
242        assert!(verify_mst(&g, &cert).is_ok());
243
244        // Tamper with the weight.
245        let mut bad = cert.clone();
246        bad.total_weight_repr = "99".to_string();
247        assert!(verify_mst(&g, &bad).is_err());
248
249        // A suboptimal spanning tree {edge1, edge2} (weight 5) is rejected by
250        // the cycle property: edge0 (weight 1) is cheaper than the path max 3.
251        let suboptimal = crate::certificate::MstCertificate {
252            edge_ids: vec![1, 2],
253            total_weight_repr: "5".to_string(),
254        };
255        assert!(verify_mst(&g, &suboptimal).is_err());
256    }
257}
sim_lib_discrete_graph/mst.rs

sim_lib_discrete_graph/
mst.rs
