Add connected_components_vec

src/algo/mod.rs

@@ -46,6 +46,21 @@ pub use k_shortest_path::k_shortest_path;
 pub use matching::{greedy_matching, maximum_matching, Matching};
 pub use simple_paths::all_simple_paths;
 
+fn connected_components_inner<G>(g: G) -> Vec<usize>
+where
+    G: NodeCompactIndexable + IntoEdgeReferences,
+{
+    let mut vertex_sets = UnionFind::new(g.node_bound());
+    for edge in g.edge_references() {
+        let (a, b) = (edge.source(), edge.target());
+
+        // union the two vertices of the edge
+        vertex_sets.union(g.to_index(a), g.to_index(b));
+    }
+
+    vertex_sets.into_labeling()
+}
+
 /// \[Generic\] Return the number of connected components of the graph.
 ///
 /// For a directed graph, this is the *weakly* connected components.
@@ -88,19 +103,73 @@ pub fn connected_components<G>(g: G) -> usize
 where
     G: NodeCompactIndexable + IntoEdgeReferences,
 {
-    let mut vertex_sets = UnionFind::new(g.node_bound());
-    for edge in g.edge_references() {
-        let (a, b) = (edge.source(), edge.target());
-
-        // union the two vertices of the edge
-        vertex_sets.union(g.to_index(a), g.to_index(b));
-    }
-    let mut labels = vertex_sets.into_labeling();
+    let mut labels = connected_components_inner(g);
     labels.sort_unstable();
     labels.dedup();
     labels.len()
 }
 
+/// \[Generic\] Return the connected components of the graph.
+///
+/// For a directed graph, this is the *weakly* connected components.
+/// # Example
+/// ```rust
+/// use petgraph::Graph;
+/// use petgraph::algo::{connected_components_vec};
+/// use petgraph::prelude::*;
+///
+/// let mut graph : Graph<(), (), Directed> = Graph::new();
+/// let a = graph.add_node(()); // node with no weight
+/// let b = graph.add_node(());
+/// let c = graph.add_node(());
+/// let d = graph.add_node(());
+/// let e = graph.add_node(());
+/// let f = graph.add_node(());
+/// let g = graph.add_node(());
+/// let h = graph.add_node(());
+///
+/// graph.extend_with_edges(&[
+///     (a, b),
+///     (b, c),
+///     (c, d),
+///     (d, a),
+///     (e, f),
+///     (f, g),
+///     (g, h),
+///     (h, e)
+/// ]);
+/// // a ----> b       e ----> f
+/// // ^       |       ^       |
+/// // |       v       |       v
+/// // d <---- c       h <---- g
+///
+/// assert_eq!(connected_components_vec(&graph), (2, vec![0, 0, 0, 0, 1, 1, 1, 1]));
+/// graph.add_edge(b,e,());
+/// assert_eq!(connected_components_vec(&graph), (1, vec![0, 0, 0, 0, 0, 0, 0, 0]));
+/// ```
+pub fn connected_components_vec<G>(g: G) -> (usize, Vec<usize>)
+where
+    G: NodeCompactIndexable + IntoEdgeReferences,
+{
+    let mut labels = connected_components_inner(g);
+
+    let mut components = HashMap::new();
+    let mut num_components = 0usize;
+    for label in labels.iter_mut() {
+        let value = components.entry(label.clone()).or_insert_with(|| {
+            let new = num_components;
+
+            num_components += 1usize;
+
+            new
+        });
+
+        *label = *value;
+    }
+
+    (num_components, labels)
+}
+
 /// \[Generic\] Return `true` if the input graph contains a cycle.
 ///
 /// Always treats the input graph as if undirected.

tests/connected_components.rs

@@ -0,0 +1,78 @@
+extern crate petgraph;
+
+use petgraph::algo::connected_components_vec;
+use petgraph::{Directed, Graph};
+
+#[test]
+fn connected_components_vec_test_simple() {
+    let mut graph: Graph<(), (), Directed> = Graph::new();
+    let a = graph.add_node(()); // node with no weight
+    let b = graph.add_node(());
+    let c = graph.add_node(());
+    let d = graph.add_node(());
+
+    let e = graph.add_node(());
+    let f = graph.add_node(());
+    let g = graph.add_node(());
+
+    let h = graph.add_node(());
+    let i = graph.add_node(());
+    let j = graph.add_node(());
+    let k = graph.add_node(());
+
+    graph.extend_with_edges(&[
+        (a, b),
+        (b, c),
+        (c, d),
+        (d, a),
+        //
+        (e, f),
+        (f, g),
+        (g, e),
+        //
+        (h, i),
+        (i, j),
+        (j, k),
+        (k, h),
+    ]);
+
+    let res = connected_components_vec(&graph);
+    assert_eq!(res.0, 3usize);
+    assert_eq!(res.1, vec![0, 0, 0, 0, 1, 1, 1, 2, 2, 2, 2]);
+}
+
+#[test]
+fn connected_components_vec_test_mixed() {
+    let mut graph: Graph<(), (), Directed> = Graph::new();
+    let a = graph.add_node(()); // node with no weight
+    let b = graph.add_node(());
+    let c = graph.add_node(());
+    let d = graph.add_node(());
+    let e = graph.add_node(());
+    let f = graph.add_node(());
+    let g = graph.add_node(());
+    let h = graph.add_node(());
+    let i = graph.add_node(());
+    let j = graph.add_node(());
+    let k = graph.add_node(());
+
+    graph.extend_with_edges(&[
+        (a, b),
+        (b, j),
+        (j, k),
+        (k, a),
+        //
+        (e, g),
+        (f, f),
+        (g, e),
+        //
+        (c, d),
+        (d, h),
+        (h, i),
+        (i, c),
+    ]);
+
+    let res = connected_components_vec(&graph);
+    assert_eq!(res.0, 4usize);
+    assert_eq!(res.1, vec![0, 0, 1, 1, 2, 3, 2, 1, 1, 0, 0]);
+}