weavatrix_graph/algo/
mst.rs

1use crate::IndexUndirectedGraphView;
2
3#[derive(Debug, Clone, PartialEq, Eq)]
4pub struct SpanningForest<Edge> {
5    edges: Vec<Edge>,
6    total_weight: u128,
7    component_count: usize,
8}
9
10impl<Edge> SpanningForest<Edge> {
11    #[must_use]
12    pub fn edges(&self) -> &[Edge] {
13        &self.edges
14    }
15
16    #[must_use]
17    pub const fn total_weight(&self) -> u128 {
18        self.total_weight
19    }
20
21    #[must_use]
22    pub const fn component_count(&self) -> usize {
23        self.component_count
24    }
25
26    #[must_use]
27    pub fn into_edges(self) -> Vec<Edge> {
28        self.edges
29    }
30}
31
32pub fn minimum_spanning_forest<G, F>(graph: &G, mut edge_weight: F) -> SpanningForest<G::Edge>
33where
34    G: IndexUndirectedGraphView,
35    F: FnMut(G::Edge) -> u64,
36{
37    let mut weighted = graph
38        .edge_indices()
39        .map(|edge| (edge_weight(edge), G::edge_slot(edge), edge))
40        .collect::<Vec<_>>();
41    weighted.sort_unstable_by_key(|&(weight, slot, _)| (weight, slot));
42
43    let mut sets = DisjointSets::new(graph.node_bound());
44    let mut selected = Vec::with_capacity(graph.node_count().saturating_sub(1));
45    let mut total_weight = 0_u128;
46    let mut component_count = graph.node_count();
47    for (weight, _, edge) in weighted {
48        let Some(endpoints) = graph.edge_endpoints(edge) else {
49            continue;
50        };
51        let source = G::node_slot(endpoints.source());
52        let target = G::node_slot(endpoints.target());
53        if sets.union(source, target) {
54            selected.push(edge);
55            total_weight += u128::from(weight);
56            component_count -= 1;
57        }
58    }
59    SpanningForest {
60        edges: selected,
61        total_weight,
62        component_count,
63    }
64}
65
66struct DisjointSets {
67    parent: Vec<usize>,
68    rank: Vec<u8>,
69}
70
71impl DisjointSets {
72    fn new(bound: usize) -> Self {
73        Self {
74            parent: (0..bound).collect(),
75            rank: vec![0; bound],
76        }
77    }
78
79    fn find(&mut self, mut node: usize) -> usize {
80        let mut root = node;
81        while self.parent[root] != root {
82            root = self.parent[root];
83        }
84        while self.parent[node] != node {
85            let parent = self.parent[node];
86            self.parent[node] = root;
87            node = parent;
88        }
89        root
90    }
91
92    fn union(&mut self, left: usize, right: usize) -> bool {
93        let mut left = self.find(left);
94        let mut right = self.find(right);
95        if left == right {
96            return false;
97        }
98        if self.rank[left] < self.rank[right] {
99            std::mem::swap(&mut left, &mut right);
100        }
101        self.parent[right] = left;
102        if self.rank[left] == self.rank[right] {
103            self.rank[left] = self.rank[left].saturating_add(1);
104        }
105        true
106    }
107}
weavatrix_graph/algo/mst.rs

weavatrix_graph/algo/
mst.rs
