weavatrix_graph/algo/
mst.rs

1use crate::IndexUndirectedGraphView;
2use crate::Vec;
3use alloc::collections::BinaryHeap;
4use core::cmp::Reverse;
5
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub struct SpanningForest<Edge> {
8    edges: Vec<Edge>,
9    total_weight: u128,
10    component_count: usize,
11}
12
13impl<Edge> SpanningForest<Edge> {
14    #[must_use]
15    pub fn edges(&self) -> &[Edge] {
16        &self.edges
17    }
18
19    #[must_use]
20    pub const fn total_weight(&self) -> u128 {
21        self.total_weight
22    }
23
24    #[must_use]
25    pub const fn component_count(&self) -> usize {
26        self.component_count
27    }
28
29    #[must_use]
30    pub fn into_edges(self) -> Vec<Edge> {
31        self.edges
32    }
33}
34
35pub fn minimum_spanning_forest<G, F>(graph: &G, mut edge_weight: F) -> SpanningForest<G::Edge>
36where
37    G: IndexUndirectedGraphView,
38    F: FnMut(G::Edge) -> u64,
39{
40    let mut weighted = graph
41        .edge_indices()
42        .map(|edge| (edge_weight(edge), G::edge_slot(edge), edge))
43        .collect::<Vec<_>>();
44    weighted.sort_unstable_by_key(|&(weight, slot, _)| (weight, slot));
45
46    let mut sets = DisjointSets::new(graph.node_bound());
47    let mut selected = Vec::with_capacity(graph.node_count().saturating_sub(1));
48    let mut total_weight = 0_u128;
49    let mut component_count = graph.node_count();
50    for (weight, _, edge) in weighted {
51        let Some(endpoints) = graph.edge_endpoints(edge) else {
52            continue;
53        };
54        let source = G::node_slot(endpoints.source());
55        let target = G::node_slot(endpoints.target());
56        if sets.union(source, target) {
57            selected.push(edge);
58            total_weight += u128::from(weight);
59            component_count -= 1;
60        }
61    }
62    SpanningForest {
63        edges: selected,
64        total_weight,
65        component_count,
66    }
67}
68
69#[must_use]
70pub fn prim_spanning_forest<G, F>(graph: &G, mut edge_weight: F) -> SpanningForest<G::Edge>
71where
72    G: IndexUndirectedGraphView,
73    F: FnMut(G::Edge) -> u64,
74{
75    let mut edges = vec![None; graph.edge_bound()];
76    let mut weights = vec![0_u64; graph.edge_bound()];
77    for edge in graph.edge_indices() {
78        let slot = G::edge_slot(edge);
79        edges[slot] = Some(edge);
80        weights[slot] = edge_weight(edge);
81    }
82    let mut seen = vec![false; graph.node_bound()];
83    let mut selected = Vec::with_capacity(graph.node_count().saturating_sub(1));
84    let mut total_weight = 0_u128;
85    let mut component_count = 0;
86    let mut queue = BinaryHeap::new();
87    for root in graph.node_indices() {
88        let root_slot = G::node_slot(root);
89        if seen[root_slot] {
90            continue;
91        }
92        component_count += 1;
93        visit_prim::<G>(graph, root, &weights, &mut seen, &mut queue);
94        while let Some(Reverse((weight, edge_slot, target_slot))) = queue.pop() {
95            if seen[target_slot] {
96                continue;
97            }
98            let Some(edge) = edges[edge_slot] else {
99                continue;
100            };
101            let Some(endpoints) = graph.edge_endpoints(edge) else {
102                continue;
103            };
104            let target = if G::node_slot(endpoints.source()) == target_slot {
105                endpoints.source()
106            } else {
107                endpoints.target()
108            };
109            if seen[G::node_slot(target)] {
110                continue;
111            }
112            selected.push(edge);
113            total_weight += u128::from(weight);
114            visit_prim::<G>(graph, target, &weights, &mut seen, &mut queue);
115        }
116    }
117    SpanningForest {
118        edges: selected,
119        total_weight,
120        component_count,
121    }
122}
123
124fn visit_prim<G>(
125    graph: &G,
126    node: G::Node,
127    weights: &[u64],
128    seen: &mut [bool],
129    queue: &mut BinaryHeap<Reverse<(u64, usize, usize)>>,
130) where
131    G: IndexUndirectedGraphView,
132{
133    seen[G::node_slot(node)] = true;
134    for edge in graph.incident_edges(node) {
135        let Some(target) = graph.opposite(edge, node) else {
136            continue;
137        };
138        let target_slot = G::node_slot(target);
139        if !seen[target_slot] {
140            let edge_slot = G::edge_slot(edge);
141            queue.push(Reverse((weights[edge_slot], edge_slot, target_slot)));
142        }
143    }
144}
145
146struct DisjointSets {
147    parent: Vec<usize>,
148    rank: Vec<u8>,
149}
150
151impl DisjointSets {
152    fn new(bound: usize) -> Self {
153        Self {
154            parent: (0..bound).collect(),
155            rank: vec![0; bound],
156        }
157    }
158
159    fn find(&mut self, mut node: usize) -> usize {
160        let mut root = node;
161        while self.parent[root] != root {
162            root = self.parent[root];
163        }
164        while self.parent[node] != node {
165            let parent = self.parent[node];
166            self.parent[node] = root;
167            node = parent;
168        }
169        root
170    }
171
172    fn union(&mut self, left: usize, right: usize) -> bool {
173        let mut left = self.find(left);
174        let mut right = self.find(right);
175        if left == right {
176            return false;
177        }
178        if self.rank[left] < self.rank[right] {
179            core::mem::swap(&mut left, &mut right);
180        }
181        self.parent[right] = left;
182        if self.rank[left] == self.rank[right] {
183            self.rank[left] = self.rank[left].saturating_add(1);
184        }
185        true
186    }
187}
weavatrix_graph/algo/mst.rs

weavatrix_graph/algo/
mst.rs
