1use 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
16type 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
29pub 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
72pub 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 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(); g.add_edge(1, 2, 2).unwrap(); g.add_edge(0, 2, 3).unwrap(); 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 let mut g: Graph<u8, u64> = Graph::with_nodes(vec![0, 1, 2, 3], Directedness::Undirected);
174 g.add_edge(0, 1, 1).unwrap(); g.add_edge(1, 2, 1).unwrap(); g.add_edge(2, 3, 1).unwrap(); g.add_edge(0, 3, 1).unwrap(); g.add_edge(0, 2, 1).unwrap(); 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 let mut bad = cert.clone();
207 bad.total_weight_repr = "99".to_string();
208 assert!(verify_mst(&g, &bad).is_err());
209
210 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}