Skip to main content

oxihuman_core/
prim_mst.rs

1// Copyright (C) 2026 COOLJAPAN OU (Team KitaSan)
2// SPDX-License-Identifier: Apache-2.0
3#![allow(dead_code)]
4
5//! Prim's minimum spanning tree algorithm.
6
7/// A weighted undirected edge in the MST.
8#[derive(Debug, Clone, Copy)]
9pub struct MstEdge {
10    pub u: usize,
11    pub v: usize,
12    pub weight: f32,
13}
14
15/// An undirected weighted graph.
16pub struct PrimGraph {
17    pub n: usize,
18    pub adj: Vec<Vec<(usize, f32)>>,
19}
20
21impl PrimGraph {
22    pub fn new(n: usize) -> Self {
23        PrimGraph {
24            n,
25            adj: vec![Vec::new(); n],
26        }
27    }
28}
29
30/// Create a new Prim graph with `n` nodes.
31pub fn new_prim_graph(n: usize) -> PrimGraph {
32    PrimGraph::new(n)
33}
34
35/// Add an undirected weighted edge.
36pub fn prim_add_edge(g: &mut PrimGraph, u: usize, v: usize, w: f32) {
37    if u < g.n && v < g.n {
38        g.adj[u].push((v, w));
39        g.adj[v].push((u, w));
40    }
41}
42
43/// Run Prim's algorithm from vertex 0. Returns the MST edges, or `None` if
44/// the graph is disconnected.
45pub fn prim_mst(g: &PrimGraph) -> Option<Vec<MstEdge>> {
46    if g.n == 0 {
47        return Some(Vec::new());
48    }
49    let mut in_mst = vec![false; g.n];
50    let mut key = vec![f32::INFINITY; g.n];
51    let mut parent = vec![usize::MAX; g.n];
52    key[0] = 0.0;
53
54    for _ in 0..g.n {
55        /* find min-key vertex not yet in MST */
56        let u = (0..g.n).filter(|&v| !in_mst[v]).min_by(|&a, &b| {
57            key[a]
58                .partial_cmp(&key[b])
59                .unwrap_or(std::cmp::Ordering::Equal)
60        })?;
61        in_mst[u] = true;
62        for &(v, w) in &g.adj[u] {
63            if !in_mst[v] && w < key[v] {
64                key[v] = w;
65                parent[v] = u;
66            }
67        }
68    }
69
70    /* collect edges */
71    let mut edges = Vec::with_capacity(g.n.saturating_sub(1));
72    for v in 1..g.n {
73        if parent[v] == usize::MAX {
74            return None; /* disconnected */
75        }
76        edges.push(MstEdge {
77            u: parent[v],
78            v,
79            weight: key[v],
80        });
81    }
82    Some(edges)
83}
84
85/// Return the total weight of the MST, or `None` if not connected.
86pub fn prim_mst_weight(g: &PrimGraph) -> Option<f32> {
87    prim_mst(g).map(|edges| edges.iter().map(|e| e.weight).sum())
88}
89
90/// Return number of nodes.
91pub fn prim_node_count(g: &PrimGraph) -> usize {
92    g.n
93}
94
95/// Return number of edges (undirected count).
96pub fn prim_edge_count(g: &PrimGraph) -> usize {
97    g.adj.iter().map(|v| v.len()).sum::<usize>() / 2
98}
99
100#[cfg(test)]
101mod tests {
102    use super::*;
103
104    #[test]
105    fn test_triangle() {
106        let mut g = new_prim_graph(3);
107        prim_add_edge(&mut g, 0, 1, 1.0);
108        prim_add_edge(&mut g, 1, 2, 2.0);
109        prim_add_edge(&mut g, 0, 2, 5.0);
110        let edges = prim_mst(&g).expect("should succeed");
111        assert_eq!(edges.len(), 2);
112        let w = prim_mst_weight(&g).expect("should succeed");
113        assert!((w - 3.0).abs() < 1e-6);
114    }
115
116    #[test]
117    fn test_single_node() {
118        let g = new_prim_graph(1);
119        let edges = prim_mst(&g).expect("should succeed");
120        assert!(edges.is_empty());
121    }
122
123    #[test]
124    fn test_empty_graph() {
125        let g = new_prim_graph(0);
126        let edges = prim_mst(&g).expect("should succeed");
127        assert!(edges.is_empty());
128    }
129
130    #[test]
131    fn test_mst_edge_count() {
132        /* n nodes -> n-1 MST edges */
133        let mut g = new_prim_graph(5);
134        prim_add_edge(&mut g, 0, 1, 1.0);
135        prim_add_edge(&mut g, 1, 2, 1.0);
136        prim_add_edge(&mut g, 2, 3, 1.0);
137        prim_add_edge(&mut g, 3, 4, 1.0);
138        let edges = prim_mst(&g).expect("should succeed");
139        assert_eq!(edges.len(), 4);
140    }
141
142    #[test]
143    fn test_minimum_weight_chosen() {
144        let mut g = new_prim_graph(3);
145        prim_add_edge(&mut g, 0, 1, 10.0);
146        prim_add_edge(&mut g, 0, 1, 1.0); /* cheaper parallel edge via adj */
147        prim_add_edge(&mut g, 1, 2, 1.0);
148        let w = prim_mst_weight(&g).expect("should succeed");
149        assert!(w <= 11.0);
150    }
151
152    #[test]
153    fn test_node_count() {
154        let g = new_prim_graph(7);
155        assert_eq!(prim_node_count(&g), 7);
156    }
157
158    #[test]
159    fn test_edge_count() {
160        let mut g = new_prim_graph(3);
161        prim_add_edge(&mut g, 0, 1, 1.0);
162        prim_add_edge(&mut g, 1, 2, 2.0);
163        assert_eq!(prim_edge_count(&g), 2);
164    }
165
166    #[test]
167    fn test_disconnected_returns_none() {
168        let mut g = new_prim_graph(4);
169        prim_add_edge(&mut g, 0, 1, 1.0);
170        /* nodes 2 and 3 are isolated */
171        assert!(prim_mst(&g).is_none());
172    }
173
174    #[test]
175    fn test_star_topology() {
176        let mut g = new_prim_graph(5);
177        for i in 1..5 {
178            prim_add_edge(&mut g, 0, i, i as f32);
179        }
180        let edges = prim_mst(&g).expect("should succeed");
181        assert_eq!(edges.len(), 4);
182    }
183}