1 //! Minimum Spanning Tree algorithms.
2
3 use std::collections::{BinaryHeap, HashMap};
4
5 use crate::prelude::*;
6
7 use crate::data::Element;
8 use crate::scored::MinScored;
9 use crate::unionfind::UnionFind;
10 use crate::visit::{Data, IntoNodeReferences, NodeRef};
11 use 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.
min_spanning_tree<G>(g: G) -> MinSpanningTree<G> where G::NodeWeight: Clone, G::EdgeWeight: Clone + PartialOrd, G: IntoNodeReferences + IntoEdgeReferences + NodeIndexable,25 pub fn min_spanning_tree<G>(g: G) -> MinSpanningTree<G>
26 where
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)]
56 pub struct MinSpanningTree<G>
57 where
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
69 impl<G> Iterator for MinSpanningTree<G>
70 where
71 G: IntoNodeReferences + NodeIndexable,
72 G::NodeWeight: Clone,
73 G::EdgeWeight: PartialOrd,
74 {
75 type Item = Element<G::NodeWeight, G::EdgeWeight>;
76
next(&mut self) -> Option<Self::Item>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 }
118