Add Prim's Algorithm for Minimum Spanning Tree (petgraph#625)

Prim's algorithm: https://en.wikipedia.org/wiki/Prim%27s_algorithm

This PR adds another algorithm to find the Minimum Spanning Tree of a
graph, an alternative to Kruskal's.

1. Prim's algorithm has some limitations, such as the input graph must
be undirected and it should not have disconnected components.
I'm wondering if this can be enforced with traits, if anyone can help on
this.

2. I would like to suggest adding another algorithm as well,
`min_spanning_forest_prim`, that iterates through this algorithm to find
min spanning tree for all input graph components.
`min_spanning_forest_kruskal` would be trivial to implement, since
current `min_spanning_tree` function, using Kruskal's algorithm, already
finds a min spanning forest. Any thoughts on this?

This PR also aims to fix current `min_spanning_tree` benches, since they
currently do not iterate through the created `MinSpanningTree`
structure, not giving the actual bench for Kruskal algorithm.

3. I added a simple function to iterate through the generated tree
elements, but I'm not very familiar with Rust ecosystem, so I would
appreciate if anyone knows a better way to force the iteration:
```Rust
// Current Bench:
bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));

// Suggested Bench:
bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));

// Force Tree Iteration:
fn iterate_mst_kruskal<G>(g: G) -> bool
where
    G: Data + IntoEdges + IntoNodeReferences + IntoEdgeReferences + NodeIndexable,
    G::NodeWeight: Clone,
    G::EdgeWeight: Clone + PartialOrd,
{
    let mst = min_spanning_tree(g);
    mst.into_iter().all(|_| true)
}
```

---------

Co-authored-by: Sambinelli <[email protected]>

Author: BryanCruz

benches/min_spanning_tree.rs

@@ -9,52 +9,101 @@ use test::Bencher;
 mod common;
 use common::{digraph, ungraph};
 
-use petgraph::algo::min_spanning_tree;
+use petgraph::{
+    algo::{min_spanning_tree, min_spanning_tree_prim},
+    visit::{Data, IntoEdgeReferences, IntoEdges, IntoNodeReferences, NodeIndexable},
+    Graph, Undirected,
+};
 
 #[bench]
-fn min_spanning_tree_praust_undir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_praust_undir_bench(bench: &mut Bencher) {
     let a = ungraph().praust_a();
     let b = ungraph().praust_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
 }
 
 #[bench]
-fn min_spanning_tree_praust_dir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_praust_dir_bench(bench: &mut Bencher) {
     let a = digraph().praust_a();
     let b = digraph().praust_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
 }
 
 #[bench]
-fn min_spanning_tree_full_undir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_full_undir_bench(bench: &mut Bencher) {
     let a = ungraph().full_a();
     let b = ungraph().full_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
 }
 
 #[bench]
-fn min_spanning_tree_full_dir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_full_dir_bench(bench: &mut Bencher) {
     let a = digraph().full_a();
     let b = digraph().full_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
 }
 
 #[bench]
-fn min_spanning_tree_petersen_undir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_petersen_undir_bench(bench: &mut Bencher) {
     let a = ungraph().petersen_a();
     let b = ungraph().petersen_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
 }
 
 #[bench]
-fn min_spanning_tree_petersen_dir_bench(bench: &mut Bencher) {
+fn min_spanning_tree_kruskal_petersen_dir_bench(bench: &mut Bencher) {
     let a = digraph().petersen_a();
     let b = digraph().petersen_b();
 
-    bench.iter(|| (min_spanning_tree(&a), min_spanning_tree(&b)));
+    bench.iter(|| (iterate_mst_kruskal(&a), iterate_mst_kruskal(&b)));
+}
+
+#[bench]
+fn min_spanning_tree_prim_praust_undir_bench(bench: &mut Bencher) {
+    let a = ungraph().praust_a();
+    let b = ungraph().praust_b();
+
+    bench.iter(|| (iterate_mst_prim(&a), iterate_mst_prim(&b)));
+}
+
+#[bench]
+fn min_spanning_tree_prim_full_undir_bench(bench: &mut Bencher) {
+    let a = ungraph().full_a();
+    let b = ungraph().full_b();
+
+    bench.iter(|| (iterate_mst_prim(&a), iterate_mst_prim(&b)));
+}
+
+#[bench]
+fn min_spanning_tree_prim_petersen_undir_bench(bench: &mut Bencher) {
+    let a: Graph<(), (), Undirected> = ungraph().petersen_a();
+    let b = ungraph().petersen_b();
+
+    bench.iter(|| (iterate_mst_prim(&a), iterate_mst_prim(&b)));
+}
+
+fn iterate_mst_kruskal<G>(g: G)
+where
+    G: Data + IntoEdges + IntoNodeReferences + IntoEdgeReferences + NodeIndexable,
+    G::NodeWeight: Clone,
+    G::EdgeWeight: Clone + PartialOrd,
+{
+    for e in min_spanning_tree(g) {
+        std::hint::black_box(e);
+    }
+}
+fn iterate_mst_prim<G>(g: G)
+where
+    G: Data + IntoEdges + IntoNodeReferences + IntoEdgeReferences + NodeIndexable,
+    G::NodeWeight: Clone,
+    G::EdgeWeight: Clone + PartialOrd,
+{
+    for e in min_spanning_tree_prim(g) {
+        std::hint::black_box(e);
+    }
 }

src/algo/min_spanning_tree.rs

@@ -1,13 +1,13 @@
 //! Minimum Spanning Tree algorithms.
 
-use std::collections::{BinaryHeap, HashMap};
+use std::collections::{BinaryHeap, HashMap, HashSet};
 
 use crate::prelude::*;
 
 use crate::data::Element;
 use crate::scored::MinScored;
 use crate::unionfind::UnionFind;
-use crate::visit::{Data, IntoNodeReferences, NodeRef};
+use crate::visit::{Data, IntoEdges, IntoNodeReferences, NodeRef};
 use crate::visit::{IntoEdgeReferences, NodeIndexable};
 
 /// \[Generic\] Compute a *minimum spanning tree* of a graph.
@@ -22,6 +22,10 @@ use crate::visit::{IntoEdgeReferences, NodeIndexable};
 /// and **|V| - c** edges, where **c** is the number of connected components in `g`.
 ///
 /// Use `from_elements` to create a graph from the resulting iterator.
+///
+/// See also: [`.min_spanning_tree_prim(g)`][1] for implementation using Prim's algorithm.
+///
+/// [1]: fn.min_spanning_tree_prim.html
 pub fn min_spanning_tree<G>(g: G) -> MinSpanningTree<G>
 where
     G::NodeWeight: Clone,
@@ -52,6 +56,8 @@ where
 }
 
 /// An iterator producing a minimum spanning forest of a graph.
+/// It will first iterate all Node elements from original graph,
+/// then iterate Edge elements from computed minimum spanning forest.
 #[derive(Debug, Clone)]
 pub struct MinSpanningTree<G>
 where
@@ -115,3 +121,138 @@ where
         None
     }
 }
+
+/// \[Generic\] Compute a *minimum spanning tree* of a graph using Prim's algorithm.
+///
+/// Graph is treated as if undirected. The computed minimum spanning tree can be wrong
+/// if this is not true.
+///
+/// Graph is treated as if connected (has only 1 component). Otherwise, the resulting
+/// graph will only contain edges for an arbitrary minimum spanning tree for a single component.
+///
+/// Using Prim's algorithm with runtime **O(|E| log |V|)**.
+///
+/// The resulting graph has all the vertices of the input graph (with identical node indices),
+/// and **|V| - 1** edges if input graph is connected, and |W| edges if disconnected, where |W| < |V| - 1.
+///
+/// Use `from_elements` to create a graph from the resulting iterator.
+///
+/// See also: [`.min_spanning_tree(g)`][1] for implementation using Kruskal's algorithm and support for minimum spanning forest.
+///
+/// [1]: fn.min_spanning_tree.html
+pub fn min_spanning_tree_prim<G>(g: G) -> MinSpanningTreePrim<G>
+where
+    G::EdgeWeight: PartialOrd,
+    G: IntoNodeReferences + IntoEdgeReferences,
+{
+    let sort_edges = BinaryHeap::with_capacity(g.edge_references().size_hint().0);
+    let nodes_taken = HashSet::with_capacity(g.node_references().size_hint().0);
+    let initial_node = g.node_references().next();
+
+    MinSpanningTreePrim {
+        graph: g,
+        node_ids: Some(g.node_references()),
+        node_map: HashMap::new(),
+        node_count: 0,
+        sort_edges,
+        nodes_taken,
+        initial_node,
+    }
+}
+
+/// An iterator producing a minimum spanning tree of a graph.
+/// It will first iterate all Node elements from original graph,
+/// then iterate Edge elements from computed minimum spanning tree.
+#[derive(Debug, Clone)]
+pub struct MinSpanningTreePrim<G>
+where
+    G: IntoNodeReferences,
+{
+    graph: G,
+    node_ids: Option<G::NodeReferences>,
+    node_map: HashMap<usize, usize>,
+    node_count: usize,
+    #[allow(clippy::type_complexity)]
+    sort_edges: BinaryHeap<MinScored<G::EdgeWeight, (G::NodeId, G::NodeId)>>,
+    nodes_taken: HashSet<usize>,
+    initial_node: Option<G::NodeRef>,
+}
+
+impl<G> Iterator for MinSpanningTreePrim<G>
+where
+    G: IntoNodeReferences + IntoEdges + NodeIndexable,
+    G::NodeWeight: Clone,
+    G::EdgeWeight: Clone + PartialOrd,
+{
+    type Item = Element<G::NodeWeight, G::EdgeWeight>;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // Iterate through Node elements
+        let g = self.graph;
+        if let Some(ref mut iter) = self.node_ids {
+            if let Some(node) = iter.next() {
+                self.node_map.insert(g.to_index(node.id()), self.node_count);
+                self.node_count += 1;
+                return Some(Element::Node {
+                    weight: node.weight().clone(),
+                });
+            }
+        }
+        self.node_ids = None;
+
+        // Bootstrap Prim's algorithm to find MST Edge elements.
+        // Mark initial node as taken and add its edges to priority queue.
+        if let Some(initial_node) = self.initial_node {
+            let initial_node_index = g.to_index(initial_node.id());
+            self.nodes_taken.insert(initial_node_index);
+
+            let initial_edges = g.edges(initial_node.id());
+            for edge in initial_edges {
+                self.sort_edges.push(MinScored(
+                    edge.weight().clone(),
+                    (edge.source(), edge.target()),
+                ));
+            }
+        };
+        self.initial_node = None;
+
+        // Clear edges queue if all nodes were already included in MST.
+        if self.nodes_taken.len() == self.node_count {
+            self.sort_edges.clear();
+        };
+
+        // Prim's algorithm:
+        // Iterate through Edge elements, adding an edge to the MST iff some of it's nodes are not part of MST yet.
+        while let Some(MinScored(score, (source, target))) = self.sort_edges.pop() {
+            let (source_index, target_index) = (g.to_index(source), g.to_index(target));
+
+            if self.nodes_taken.contains(&target_index) {
+                continue;
+            }
+
+            self.nodes_taken.insert(target_index);
+            for edge in g.edges(target) {
+                self.sort_edges.push(MinScored(
+                    edge.weight().clone(),
+                    (edge.source(), edge.target()),
+                ));
+            }
+
+            let (&source_order, &target_order) = match (
+                self.node_map.get(&source_index),
+                self.node_map.get(&target_index),
+            ) {
+                (Some(source_order), Some(target_order)) => (source_order, target_order),
+                _ => panic!("Edge references unknown node"),
+            };
+
+            return Some(Element::Edge {
+                source: source_order,
+                target: target_order,
+                weight: score,
+            });
+        }
+
+        None
+    }
+}

src/algo/mod.rs

@@ -47,7 +47,7 @@ pub use isomorphism::{
 };
 pub use k_shortest_path::k_shortest_path;
 pub use matching::{greedy_matching, maximum_matching, Matching};
-pub use min_spanning_tree::min_spanning_tree;
+pub use min_spanning_tree::{min_spanning_tree, min_spanning_tree_prim};
 pub use page_rank::page_rank;
 pub use simple_paths::all_simple_paths;