Merge branch 'main' into remove-39-retworkx

Author: IvanIsCoding

releasenotes/notes/add-bfs-layers-91bb3773c37a027d.yaml

@@ -0,0 +1,22 @@
+
+---
+features:
+  - |
+    Added a new :func:`~rustworkx.bfs_layers` (and it's per type variants
+    :func:`~rustworkx.graph_bfs_layers` and :func:`~rustworkx.digraph_bfs_layers`)
+    that performs a breadth-first search traversal and returns the nodes organized
+    by their BFS layers/levels. Each layer contains all nodes at the same distance
+    from the source nodes. This is useful for analyzing graph structure and 
+    implementing algorithms that need to process nodes level by level.
+    For example:
+
+    .. jupyter-execute::
+
+      import rustworkx
+
+      graph = rustworkx.PyDiGraph()
+      graph.extend_from_edge_list([(0, 1), (0, 2), (1, 3), (2, 3), (3, 4)])
+      
+      # Print the layers of the graph in BFS-order relative to 0 (source)
+      layers = rustworkx.bfs_layers(graph, sources=[0])
+      print('BFS layers:', layers)

rustworkx-core/src/traversal/bfs_visit.rs

@@ -10,11 +10,10 @@
 // License for the specific language governing permissions and limitations
 // under the License.
 
+use super::try_control;
 use petgraph::visit::{ControlFlow, EdgeRef, IntoEdges, VisitMap, Visitable};
 use std::collections::VecDeque;
 
-use super::try_control;
-
 /// A breadth first search (BFS) visitor event.
 #[derive(Copy, Clone, Debug)]
 pub enum BfsEvent<N, E> {
@@ -248,3 +247,37 @@ where
 
     C::continuing()
 }
+
+pub fn bfs_layers<G, I>(graph: G, sources: I) -> Vec<Vec<G::NodeId>>
+where
+    G: IntoEdges + Visitable,
+    I: IntoIterator<Item = G::NodeId>,
+    G::NodeId: Copy + std::hash::Hash + Eq,
+{
+    let mut visited = hashbrown::HashSet::new();
+    let mut current_layer: Vec<G::NodeId> = sources.into_iter().collect();
+
+    for &node in &current_layer {
+        visited.insert(node);
+    }
+
+    let mut layers: Vec<Vec<G::NodeId>> = Vec::new();
+
+    while !current_layer.is_empty() {
+        layers.push(current_layer.clone());
+
+        let mut next_layer = Vec::new();
+        for &node in &current_layer {
+            for edge in graph.edges(node) {
+                let child = edge.target();
+                if visited.insert(child) {
+                    next_layer.push(child);
+                }
+            }
+        }
+
+        current_layer = next_layer;
+    }
+
+    layers
+}

rustworkx-core/src/traversal/mod.rs

@@ -24,7 +24,7 @@ use petgraph::visit::Reversed;
 use petgraph::visit::VisitMap;
 use petgraph::visit::Visitable;
 
-pub use bfs_visit::{breadth_first_search, BfsEvent};
+pub use bfs_visit::{bfs_layers, breadth_first_search, BfsEvent};
 pub use dfs_edges::dfs_edges;
 pub use dfs_visit::{depth_first_search, DfsEvent};
 pub use dijkstra_visit::{dijkstra_search, DijkstraEvent};

rustworkx/__init__.py

@@ -2311,3 +2311,21 @@ def write_graphml(graph, path, /, keys=None, compression=None):
     :raises RuntimeError: when an error is encountered while writing the GraphML file.
     """
     raise TypeError(f"Invalid Input Type {type(graph)} for graph")
+
+
+@_rustworkx_dispatch
+def bfs_layers(graph, sources=None):
+    """Return the BFS layers of a graph as a list of lists.
+
+    :param graph: The input graph to use. Can either be a
+        :class:`~rustworkx.PyGraph` or :class:`~rustworkx.PyDiGraph`
+    :param sources: An optional list of node indices to use as the starting
+        nodes for the BFS traversal. If not specified, all nodes in the graph
+        will be used as sources.
+    :type sources: list[int]
+
+    :returns: A list of lists where each inner list contains the node indices
+        at that BFS layer/level from the source nodes
+    :rtype: list[list[int]]
+    """
+    raise TypeError(f"Invalid Input Type {type(graph)} for graph")

rustworkx/__init__.pyi

@@ -255,6 +255,8 @@ from .rustworkx import digraph_transitivity as digraph_transitivity
 from .rustworkx import graph_transitivity as graph_transitivity
 from .rustworkx import digraph_bfs_search as digraph_bfs_search
 from .rustworkx import graph_bfs_search as graph_bfs_search
+from .rustworkx import digraph_bfs_layers as digraph_bfs_layers
+from .rustworkx import graph_bfs_layers as graph_bfs_layers
 from .rustworkx import digraph_dfs_search as digraph_dfs_search
 from .rustworkx import graph_dfs_search as graph_dfs_search
 from .rustworkx import digraph_dijkstra_search as digraph_dijkstra_search
@@ -674,3 +676,7 @@ def write_graphml(
     keys: list[GraphMLKey] | None = ...,
     compression: str | None = ...,
 ) -> None: ...
+def bfs_layers(
+    graph: PyGraph | PyDiGraph,
+    sources: Sequence[int] | None = ...,
+) -> list[list[int]]: ...

rustworkx/rustworkx.pyi

@@ -1082,6 +1082,14 @@ def graph_dfs_search(
     source: Sequence[int] | None = ...,
     visitor: _DFSVisitor | None = ...,
 ) -> None: ...
+def digraph_bfs_layers(
+    digraph: PyDiGraph,
+    sources: Sequence[int] | None = ...,
+) -> list[list[int]]: ...
+def graph_bfs_layers(
+    graph: PyGraph,
+    sources: Sequence[int] | None = ...,
+) -> list[list[int]]: ...
 def digraph_dijkstra_search(
     graph: PyDiGraph,
     source: Sequence[int] | None = ...,

src/lib.rs

@@ -664,6 +664,8 @@ fn rustworkx(py: Python<'_>, m: &Bound<PyModule>) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(steiner_tree::steiner_tree))?;
     m.add_wrapped(wrap_pyfunction!(digraph_dfs_search))?;
     m.add_wrapped(wrap_pyfunction!(graph_dfs_search))?;
+    m.add_wrapped(wrap_pyfunction!(digraph_bfs_layers))?;
+    m.add_wrapped(wrap_pyfunction!(graph_bfs_layers))?;
     m.add_wrapped(wrap_pyfunction!(articulation_points))?;
     m.add_wrapped(wrap_pyfunction!(bridges))?;
     m.add_wrapped(wrap_pyfunction!(biconnected_components))?;

src/traversal/mod.rs

@@ -19,7 +19,7 @@ use dfs_visit::{dfs_handler, PyDfsVisitor};
 use dijkstra_visit::{dijkstra_handler, PyDijkstraVisitor};
 
 use rustworkx_core::traversal::{
-    ancestors as core_ancestors, bfs_predecessors as core_bfs_predecessors,
+    ancestors as core_ancestors, bfs_layers, bfs_predecessors as core_bfs_predecessors,
     bfs_successors as core_bfs_successors, breadth_first_search, depth_first_search,
     descendants as core_descendants, dfs_edges, dijkstra_search,
 };
@@ -32,6 +32,8 @@ use hashbrown::HashSet;
 
 use pyo3::exceptions::{PyIndexError, PyTypeError};
 use pyo3::prelude::*;
+use pyo3::types::PyList;
+use pyo3::IntoPyObjectExt;
 use pyo3::Python;
 
 use petgraph::graph::NodeIndex;
@@ -1116,3 +1118,77 @@ pub fn graph_dijkstra_search(
 
     Ok(())
 }
+
+/// Return the BFS layers of a PyGraph as a list of lists.
+///
+/// :param graph: The input PyGraph to use for BFS traversal
+/// :type graph: PyGraph
+/// :param sources: An optional list of node indices to use as the starting
+///     nodes for the BFS traversal. If not specified, all nodes in the graph
+///     will be used as sources.
+/// :type sources: list[int] or None
+///
+/// :returns: A list of lists where each inner list contains the node indices
+///     at that BFS layer/level from the source nodes
+/// :rtype: list[list[int]]
+#[pyfunction]
+#[pyo3(signature = (graph, sources=None))]
+pub fn graph_bfs_layers(
+    py: Python,
+    graph: &graph::PyGraph,
+    sources: Option<Vec<usize>>,
+) -> PyResult<PyObject> {
+    let starts: Vec<NodeIndex> = match sources {
+        Some(v) => v.into_iter().map(NodeIndex::new).collect(),
+        None => graph.graph.node_indices().collect(),
+    };
+
+    validate_source_nodes(&graph.graph, &starts)?;
+
+    let layers = bfs_layers(&graph.graph, starts);
+
+    let py_layers = PyList::empty(py);
+    for layer in layers {
+        let ids: Vec<usize> = layer.into_iter().map(|n| n.index()).collect();
+        let sublist = PyList::new(py, &ids)?;
+        py_layers.append(sublist)?;
+    }
+    py_layers.into_py_any(py)
+}
+
+/// Return the BFS layers of a PyDiGraph as a list of lists.
+///
+/// :param graph: The input PyDiGraph to use for BFS traversal
+/// :type graph: PyDiGraph
+/// :param sources: An optional list of node indices to use as the starting
+///     nodes for the BFS traversal. If not specified, all nodes in the graph
+///     will be used as sources.
+/// :type sources: list[int] or None
+///
+/// :returns: A list of lists where each inner list contains the node indices
+///     at that BFS layer/level from the source nodes
+/// :rtype: list[list[int]]
+#[pyfunction]
+#[pyo3(signature = (digraph, sources=None))]
+pub fn digraph_bfs_layers(
+    py: Python,
+    digraph: &digraph::PyDiGraph,
+    sources: Option<Vec<usize>>,
+) -> PyResult<PyObject> {
+    let starts: Vec<NodeIndex> = match sources {
+        Some(v) => v.into_iter().map(NodeIndex::new).collect(),
+        None => digraph.graph.node_indices().collect(),
+    };
+
+    validate_source_nodes(&digraph.graph, &starts)?;
+
+    let layers = bfs_layers(&digraph.graph, starts);
+
+    let py_layers = PyList::empty(py);
+    for layer in layers {
+        let ids: Vec<usize> = layer.into_iter().map(|n| n.index()).collect();
+        let sublist = PyList::new(py, &ids)?;
+        py_layers.append(sublist)?;
+    }
+    py_layers.into_py_any(py)
+}

tests/digraph/test_bfs_layer.py

@@ -0,0 +1,29 @@
+import unittest
+import rustworkx
+
+
+class TestDiGraphBfsLayers(unittest.TestCase):
+    def setUp(self):
+        self.graph = rustworkx.generators.path_graph(6).to_directed()
+
+    def test_simple_chain(self):
+        layers = rustworkx.bfs_layers(self.graph, [3])
+        self.assertEqual([sorted(layer) for layer in layers], [[3], [2, 4], [1, 5], [0]])
+
+    def test_multiple_sources(self):
+        layers = rustworkx.bfs_layers(self.graph, [0, 3])
+        self.assertEqual(sorted(layers[0]), [0, 3])
+
+    def test_disconnected_digraph(self):
+        g = rustworkx.PyDiGraph()
+        g.extend_from_edge_list([(0, 1), (2, 3)])
+        layers = rustworkx.bfs_layers(g, [2])
+        self.assertEqual(layers, [[2], [3]])
+
+    def test_no_sources_defaults(self):
+        layers = rustworkx.bfs_layers(self.graph, None)
+        self.assertTrue(any(0 in layer for layer in layers))
+
+    def test_invalid_source(self):
+        with self.assertRaises(IndexError):
+            rustworkx.bfs_layers(self.graph, [42])

tests/graph/test_bfs_layer.py

@@ -0,0 +1,29 @@
+import unittest
+import rustworkx
+
+
+class TestGraphBfsLayers(unittest.TestCase):
+    def setUp(self):
+        self.graph = rustworkx.generators.path_graph(5)
+
+    def test_simple_path(self):
+        layers = rustworkx.bfs_layers(self.graph, [0])
+        self.assertEqual(layers, [[0], [1], [2], [3], [4]])
+
+    def test_multiple_sources(self):
+        layers = rustworkx.bfs_layers(self.graph, [0, 4])
+        self.assertEqual(layers, [[0, 4], [1, 3], [2]])
+
+    def test_disconnected_graph(self):
+        g = rustworkx.PyGraph()
+        g.extend_from_edge_list([(0, 1), (2, 3)])
+        layers = rustworkx.bfs_layers(g, [0])
+        self.assertEqual(layers, [[0], [1]])
+
+    def test_no_sources_default_all_nodes(self):
+        layers = rustworkx.bfs_layers(self.graph, None)
+        self.assertTrue(all(isinstance(layer, list) for layer in layers))
+
+    def test_invalid_source(self):
+        with self.assertRaises(IndexError):
+            rustworkx.bfs_layers(self.graph, [99])