petgraph/algo/
min_spanning_tree.rs

1//! Minimum Spanning Tree algorithms.
2
3use std::collections::{BinaryHeap, HashMap};
4
5use crate::prelude::*;
6
7use crate::data::Element;
8use crate::scored::MinScored;
9use crate::unionfind::UnionFind;
10use crate::visit::{Data, IntoNodeReferences, NodeRef};
11use crate::visit::{IntoEdgeReferences, NodeIndexable};
12
13/// \[Generic\] Compute a *minimum spanning tree* of a graph.
14///
15/// The input graph is treated as if undirected.
16///
17/// Using Kruskal's algorithm with runtime **O(|E| log |E|)**. We actually
18/// return a minimum spanning forest, i.e. a minimum spanning tree for each connected
19/// component of the graph.
20///
21/// The resulting graph has all the vertices of the input graph (with identical node indices),
22/// and **|V| - c** edges, where **c** is the number of connected components in `g`.
23///
24/// Use `from_elements` to create a graph from the resulting iterator.
25pub fn min_spanning_tree<G>(g: G) -> MinSpanningTree<G>
26where
27    G::NodeWeight: Clone,
28    G::EdgeWeight: Clone + PartialOrd,
29    G: IntoNodeReferences + IntoEdgeReferences + NodeIndexable,
30{
31    // Initially each vertex is its own disjoint subgraph, track the connectedness
32    // of the pre-MST with a union & find datastructure.
33    let subgraphs = UnionFind::new(g.node_bound());
34
35    let edges = g.edge_references();
36    let mut sort_edges = BinaryHeap::with_capacity(edges.size_hint().0);
37    for edge in edges {
38        sort_edges.push(MinScored(
39            edge.weight().clone(),
40            (edge.source(), edge.target()),
41        ));
42    }
43
44    MinSpanningTree {
45        graph: g,
46        node_ids: Some(g.node_references()),
47        subgraphs,
48        sort_edges,
49        node_map: HashMap::new(),
50        node_count: 0,
51    }
52}
53
54/// An iterator producing a minimum spanning forest of a graph.
55#[derive(Debug, Clone)]
56pub struct MinSpanningTree<G>
57where
58    G: Data + IntoNodeReferences,
59{
60    graph: G,
61    node_ids: Option<G::NodeReferences>,
62    subgraphs: UnionFind<usize>,
63    #[allow(clippy::type_complexity)]
64    sort_edges: BinaryHeap<MinScored<G::EdgeWeight, (G::NodeId, G::NodeId)>>,
65    node_map: HashMap<usize, usize>,
66    node_count: usize,
67}
68
69impl<G> Iterator for MinSpanningTree<G>
70where
71    G: IntoNodeReferences + NodeIndexable,
72    G::NodeWeight: Clone,
73    G::EdgeWeight: PartialOrd,
74{
75    type Item = Element<G::NodeWeight, G::EdgeWeight>;
76
77    fn next(&mut self) -> Option<Self::Item> {
78        let g = self.graph;
79        if let Some(ref mut iter) = self.node_ids {
80            if let Some(node) = iter.next() {
81                self.node_map.insert(g.to_index(node.id()), self.node_count);
82                self.node_count += 1;
83                return Some(Element::Node {
84                    weight: node.weight().clone(),
85                });
86            }
87        }
88        self.node_ids = None;
89
90        // Kruskal's algorithm.
91        // Algorithm is this:
92        //
93        // 1. Create a pre-MST with all the vertices and no edges.
94        // 2. Repeat:
95        //
96        //  a. Remove the shortest edge from the original graph.
97        //  b. If the edge connects two disjoint trees in the pre-MST,
98        //     add the edge.
99        while let Some(MinScored(score, (a, b))) = self.sort_edges.pop() {
100            // check if the edge would connect two disjoint parts
101            let (a_index, b_index) = (g.to_index(a), g.to_index(b));
102            if self.subgraphs.union(a_index, b_index) {
103                let (&a_order, &b_order) =
104                    match (self.node_map.get(&a_index), self.node_map.get(&b_index)) {
105                        (Some(a_id), Some(b_id)) => (a_id, b_id),
106                        _ => panic!("Edge references unknown node"),
107                    };
108                return Some(Element::Edge {
109                    source: a_order,
110                    target: b_order,
111                    weight: score,
112                });
113            }
114        }
115        None
116    }
117}
petgraph/algo/min_spanning_tree.rs

petgraph/algo/
min_spanning_tree.rs
