New method subgraph_with_nodemap (#1461)

* adding initial subgraph_with_nodemap implementation

* adding digraph with nodemap

* fix style

* fixing return type annotation

* adding subgraph reference sphinx note

* add release note

* trying to fix release notes

Author: JoOkuma

releasenotes/notes/subgraph-with-nodemap-a6ed3bb9bbbe93cd.yaml

@@ -0,0 +1,7 @@
+---
+features:
+  - |
+    Added new :meth:`~.PyGraph.subgraph_with_nodemap` and :meth:`~.PyDiGraph.subgraph_with_nodemap`
+    methods to the :class:`~.PyGraph` and :class:`~.PyDiGraph` classes. These methods extend the
+    existing :meth:`~.PyGraph.subgraph` method by returning a :class:`~.NodeMap` object that maps
+    the nodes of the subgraph to the nodes of the original graph.

rustworkx/rustworkx.pyi

@@ -1320,6 +1320,9 @@ class PyGraph(Generic[_S, _T]):
     def remove_node(self, node: int, /) -> None: ...
     def remove_nodes_from(self, index_list: Iterable[int], /) -> None: ...
     def subgraph(self, nodes: Sequence[int], /, preserve_attrs: bool = ...) -> PyGraph[_S, _T]: ...
+    def subgraph_with_nodemap(
+        self, nodes: Sequence[int], /, preserve_attrs: bool = ...
+    ) -> tuple[PyGraph[_S, _T], NodeMap]: ...
     def substitute_node_with_subgraph(
         self,
         node: int,
@@ -1528,6 +1531,9 @@ class PyDiGraph(Generic[_S, _T]):
     def subgraph(
         self, nodes: Sequence[int], /, preserve_attrs: bool = ...
     ) -> PyDiGraph[_S, _T]: ...
+    def subgraph_with_nodemap(
+        self, nodes: Sequence[int], /, preserve_attrs: bool = ...
+    ) -> tuple[PyDiGraph[_S, _T], NodeMap]: ...
     def substitute_node_with_subgraph(
         self,
         node: int,

src/digraph.rs

@@ -3013,7 +3013,13 @@ impl PyDiGraph {
         Ok(res.index())
     }
 
-    /// Return a new PyDiGraph object for a subgraph of this graph
+    /// Return a new PyDiGraph object for a subgraph of this graph and a NodeMap
+    /// object that maps the nodes of the subgraph to the nodes of the original graph.
+    ///
+    /// .. note::
+    ///     This method is identical to :meth:`.subgraph()` but includes a
+    ///     NodeMap object that maps the nodes of the subgraph to the nodes of
+    ///     the original graph.
     ///
     /// :param list[int] nodes: A list of node indices to generate the subgraph
     ///     from. If a node index is included that is not present in the graph
@@ -3023,24 +3029,33 @@ impl PyDiGraph {
     ///     subgraph. By default this is set to ``False`` and the :attr:`~.PyDiGraph.attrs`
     ///     attribute will be ``None`` in the subgraph.
     ///
-    /// :returns: A new PyDiGraph object representing a subgraph of this graph.
+    /// :returns: A tuple containing a new PyDiGraph object representing a subgraph of this graph
+    ///     and a NodeMap object that maps the nodes of the subgraph to the nodes of the original graph.
     ///     It is worth noting that node and edge weight/data payloads are
     ///     passed by reference so if you update (not replace) an object used
     ///     as the weight in graph or the subgraph it will also be updated in
-    ///     the other. Node and edge the indices will be recreated for the subgraph for compactness.
-    ///     Therefore, do not access data using the original graph's indices.
-    /// :rtype: PyGraph
+    ///     the other.
+    /// :rtype: tuple[PyDiGraph, NodeMap]
     ///
-    #[pyo3(signature=(nodes, preserve_attrs=false),text_signature = "(self, nodes, /, preserve_attrs=False)")]
-    pub fn subgraph(&self, py: Python, nodes: Vec<usize>, preserve_attrs: bool) -> PyDiGraph {
+    #[pyo3(signature=(nodes, preserve_attrs=false), text_signature = "(self, nodes, /, preserve_attrs=False)")]
+    pub fn subgraph_with_nodemap(
+        &self,
+        py: Python,
+        nodes: Vec<usize>,
+        preserve_attrs: bool,
+    ) -> (PyDiGraph, NodeMap) {
         let node_set: HashSet<usize> = nodes.iter().cloned().collect();
+        // mapping from original node index to new node index
         let mut node_map: HashMap<NodeIndex, NodeIndex> = HashMap::with_capacity(nodes.len());
+        // mapping from new node index to original node index
+        let mut node_dict: DictMap<usize, usize> = DictMap::with_capacity(nodes.len());
         let node_filter = |node: NodeIndex| -> bool { node_set.contains(&node.index()) };
         let mut out_graph = StablePyGraph::<Directed>::new();
         let filtered = NodeFiltered(&self.graph, node_filter);
         for node in filtered.node_references() {
             let new_node = out_graph.add_node(node.1.clone_ref(py));
             node_map.insert(node.0, new_node);
+            node_dict.insert(new_node.index(), node.0.index());
         }
         for edge in filtered.edge_references() {
             let new_source = *node_map.get(&edge.source()).unwrap();
@@ -3052,14 +3067,45 @@ impl PyDiGraph {
         } else {
             py.None()
         };
-        PyDiGraph {
+        let node_map = NodeMap {
+            node_map: node_dict,
+        };
+        let subgraph = PyDiGraph {
             graph: out_graph,
             node_removed: false,
             cycle_state: algo::DfsSpace::default(),
             check_cycle: self.check_cycle,
             multigraph: self.multigraph,
             attrs,
-        }
+        };
+        (subgraph, node_map)
+    }
+
+    /// Return a new PyDiGraph object for a subgraph of this graph.
+    ///
+    /// .. note::
+    ///     To return a NodeMap object that maps the nodes of the subgraph to
+    ///     the nodes of the original graph, use :meth:`.subgraph_with_nodemap()`.
+    ///
+    /// :param list[int] nodes: A list of node indices to generate the subgraph
+    ///     from. If a node index is included that is not present in the graph
+    ///     it will silently be ignored.
+    /// :param bool preserve_attrs: If set to the True the attributes of the PyDiGraph
+    ///     will be copied by reference to be the attributes of the output
+    ///     subgraph. By default this is set to False and the :attr:`~.PyDiGraph.attrs`
+    ///     attribute will be ``None`` in the subgraph.
+    ///
+    /// :returns: A new PyDiGraph object representing a subgraph of this graph.
+    ///     It is worth noting that node and edge weight/data payloads are
+    ///     passed by reference so if you update (not replace) an object used
+    ///     as the weight in graph or the subgraph it will also be updated in
+    ///     the other.
+    /// :rtype: PyDiGraph
+    ///
+    #[pyo3(signature=(nodes, preserve_attrs=false), text_signature = "(self, nodes, /, preserve_attrs=False)")]
+    pub fn subgraph(&self, py: Python, nodes: Vec<usize>, preserve_attrs: bool) -> PyDiGraph {
+        let (subgraph, _) = self.subgraph_with_nodemap(py, nodes, preserve_attrs);
+        subgraph
     }
 
     /// Return a new PyDiGraph object for an edge induced subgraph of this graph

src/graph.rs

@@ -1907,7 +1907,13 @@ impl PyGraph {
         Ok(res.index())
     }
 
-    /// Return a new PyGraph object for a subgraph of this graph
+    /// Return a new PyGraph object for a subgraph of this graph and a NodeMap
+    /// object that maps the nodes of the subgraph to the nodes of the original graph.
+    ///
+    /// .. note::
+    ///     This method is identical to :meth:`.subgraph()` but includes a
+    ///     NodeMap object that maps the nodes of the subgraph to the nodes of
+    ///     the original graph.
     ///
     /// :param list[int] nodes: A list of node indices to generate the subgraph
     ///     from. If a node index is included that is not present in the graph
@@ -1917,23 +1923,33 @@ impl PyGraph {
     ///     subgraph. By default this is set to False and the :attr:`~.PyGraph.attrs`
     ///     attribute will be ``None`` in the subgraph.
     ///
-    /// :returns: A new PyGraph object representing a subgraph of this graph.
+    /// :returns: A tuple containing a new PyGraph object representing a subgraph of this graph
+    ///     and a NodeMap object that maps the nodes of the subgraph to the nodes of the original graph.
     ///     It is worth noting that node and edge weight/data payloads are
     ///     passed by reference so if you update (not replace) an object used
     ///     as the weight in graph or the subgraph it will also be updated in
     ///     the other.
-    /// :rtype: PyGraph
+    /// :rtype: tuple[PyGraph, NodeMap]
     ///
     #[pyo3(signature=(nodes, preserve_attrs=false), text_signature = "(self, nodes, /, preserve_attrs=False)")]
-    pub fn subgraph(&self, py: Python, nodes: Vec<usize>, preserve_attrs: bool) -> PyGraph {
+    pub fn subgraph_with_nodemap(
+        &self,
+        py: Python,
+        nodes: Vec<usize>,
+        preserve_attrs: bool,
+    ) -> (PyGraph, NodeMap) {
         let node_set: HashSet<usize> = nodes.iter().cloned().collect();
+        // mapping from original node index to new node index
         let mut node_map: HashMap<NodeIndex, NodeIndex> = HashMap::with_capacity(nodes.len());
+        // mapping from new node index to original node index
+        let mut node_dict: DictMap<usize, usize> = DictMap::with_capacity(nodes.len());
         let node_filter = |node: NodeIndex| -> bool { node_set.contains(&node.index()) };
         let mut out_graph = StablePyGraph::<Undirected>::default();
         let filtered = NodeFiltered(&self.graph, node_filter);
         for node in filtered.node_references() {
             let new_node = out_graph.add_node(node.1.clone_ref(py));
             node_map.insert(node.0, new_node);
+            node_dict.insert(new_node.index(), node.0.index());
         }
         for edge in filtered.edge_references() {
             let new_source = *node_map.get(&edge.source()).unwrap();
@@ -1945,12 +1961,43 @@ impl PyGraph {
         } else {
             py.None()
         };
-        PyGraph {
+        let node_map = NodeMap {
+            node_map: node_dict,
+        };
+        let subgraph = PyGraph {
             graph: out_graph,
             node_removed: false,
             multigraph: self.multigraph,
             attrs,
-        }
+        };
+        (subgraph, node_map)
+    }
+
+    /// Return a new PyGraph object for a subgraph of this graph.
+    ///
+    /// .. note::
+    ///     To return a NodeMap object that maps the nodes of the subgraph to
+    ///     the nodes of the original graph, use :meth:`.subgraph_with_nodemap()`.
+    ///
+    /// :param list[int] nodes: A list of node indices to generate the subgraph
+    ///     from. If a node index is included that is not present in the graph
+    ///     it will silently be ignored.
+    /// :param bool preserve_attrs: If set to the True the attributes of the PyGraph
+    ///     will be copied by reference to be the attributes of the output
+    ///     subgraph. By default this is set to False and the :attr:`~.PyGraph.attrs`
+    ///     attribute will be ``None`` in the subgraph.
+    ///
+    /// :returns: A new PyGraph object representing a subgraph of this graph.
+    ///     It is worth noting that node and edge weight/data payloads are
+    ///     passed by reference so if you update (not replace) an object used
+    ///     as the weight in graph or the subgraph it will also be updated in
+    ///     the other.
+    /// :rtype: PyGraph
+    ///
+    #[pyo3(signature=(nodes, preserve_attrs=false), text_signature = "(self, nodes, /, preserve_attrs=False)")]
+    pub fn subgraph(&self, py: Python, nodes: Vec<usize>, preserve_attrs: bool) -> PyGraph {
+        let (subgraph, _) = self.subgraph_with_nodemap(py, nodes, preserve_attrs);
+        subgraph
     }
 
     /// Return a new PyGraph object for an edge induced subgraph of this graph

tests/digraph/test_subgraph.py

@@ -138,7 +138,7 @@ def test_edge_subgraph_non_edge(self):
         self.assertEqual([(0, 1, 2), (0, 1, 3), (1, 2, 4)], subgraph.weighted_edge_list())
 
     def test_preserve_attrs(self):
-        graph = rustworkx.PyGraph(attrs="My attribute")
+        graph = rustworkx.PyDiGraph(attrs="My attribute")
         graph.add_node("a")
         graph.add_node("b")
         graph.add_node("c")
@@ -148,3 +148,62 @@ def test_preserve_attrs(self):
         self.assertEqual([(0, 1, 4)], subgraph.weighted_edge_list())
         self.assertEqual(["b", "d"], subgraph.nodes())
         self.assertEqual(graph.attrs, subgraph.attrs)
+
+    def test_subgraph_with_nodemap(self):
+        graph = rustworkx.PyDiGraph()
+        graph.add_nodes_from(list(range(6)))
+        graph.add_edges_from([(0, 1, 1), (1, 2, 2), (2, 3, 3), (3, 4, 4), (4, 5, 5)])
+
+        # Test basic subgraph with node mapping
+        subgraph, node_map = graph.subgraph_with_nodemap([0, 2, 4])
+        self.assertEqual([], subgraph.weighted_edge_list())  # No edges between disconnected nodes
+        self.assertEqual([0, 2, 4], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 0, 1: 2, 2: 4})
+
+        # Test with connected nodes
+        subgraph2, node_map2 = graph.subgraph_with_nodemap([1, 2, 3])
+        self.assertEqual([(0, 1, 2), (1, 2, 3)], subgraph2.weighted_edge_list())
+        self.assertEqual([1, 2, 3], subgraph2.nodes())
+        self.assertEqual(dict(node_map2), {0: 1, 1: 2, 2: 3})
+
+    def test_subgraph_with_nodemap_edge_cases(self):
+        graph = rustworkx.PyDiGraph()
+        graph.add_nodes_from(["a", "b", "c"])
+        graph.add_edges_from([(0, 1, 1), (1, 2, 2)])
+
+        # Test empty node list
+        subgraph, node_map = graph.subgraph_with_nodemap([])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(0, len(subgraph))
+        self.assertEqual(dict(node_map), {})
+
+        # Test invalid node indices (should be silently ignored)
+        subgraph, node_map = graph.subgraph_with_nodemap([42, 100])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(0, len(subgraph))
+        self.assertEqual(dict(node_map), {})
+
+        # Test single node (no edges in subgraph)
+        subgraph, node_map = graph.subgraph_with_nodemap([1])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(["b"], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 1})
+
+        # Test all nodes
+        subgraph, node_map = graph.subgraph_with_nodemap([0, 1, 2])
+        self.assertEqual([(0, 1, 1), (1, 2, 2)], subgraph.weighted_edge_list())
+        self.assertEqual(["a", "b", "c"], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 0, 1: 1, 2: 2})
+
+    def test_subgraph_with_nodemap_preserve_attrs(self):
+        graph = rustworkx.PyDiGraph(attrs="test_attrs")
+        graph.add_nodes_from(["a", "b", "c"])
+        graph.add_edges_from([(0, 1, 1), (1, 2, 2)])
+
+        # Test preserve_attrs=False (default)
+        subgraph, node_map = graph.subgraph_with_nodemap([0, 1])
+        self.assertIsNone(subgraph.attrs)
+
+        # Test preserve_attrs=True
+        subgraph2, node_map2 = graph.subgraph_with_nodemap([0, 1], preserve_attrs=True)
+        self.assertEqual(graph.attrs, subgraph2.attrs)

tests/graph/test_subgraph.py

@@ -148,3 +148,63 @@ def test_preserve_attrs(self):
         self.assertEqual([(0, 1, 4)], subgraph.weighted_edge_list())
         self.assertEqual(["b", "d"], subgraph.nodes())
         self.assertEqual(graph.attrs, subgraph.attrs)
+
+    def test_subgraph_with_nodemap(self):
+        graph = rustworkx.PyGraph()
+        graph.add_nodes_from(list(range(6)))
+        # Create a more complex graph: 0-1-2-3-4-5 with additional edges 1-4, 2-5
+        graph.extend_from_edge_list([(0, 1), (1, 2), (2, 3), (3, 4), (4, 5), (1, 4), (2, 5)])
+
+        # Test subset with multiple edges
+        subgraph, node_map = graph.subgraph_with_nodemap([1, 2, 4])
+        self.assertEqual(set(subgraph.node_indices()), {0, 1, 2})
+
+        # Check that we have the correct edges: 1-2 and 1-4 from original
+        edge_list = list(subgraph.edge_list())
+        self.assertEqual(len(edge_list), 2)
+        self.assertIn((0, 1), edge_list)  # 1-2 in original
+        self.assertIn((0, 2), edge_list)  # 1-4 in original
+
+        # Verify node mapping
+        node_map_dict = dict(node_map)
+        self.assertEqual(len(node_map_dict), 3)
+        self.assertEqual(set(node_map_dict.values()), {1, 2, 4})
+
+        # Test disconnected nodes (nodes with no edges between them)
+        graph2 = rustworkx.PyGraph()
+        graph2.add_nodes_from(["a", "b", "c", "d", "e"])
+        graph2.add_edges_from([(0, 1, 1), (2, 3, 2)])  # Two separate components
+        subgraph, node_map = graph2.subgraph_with_nodemap([0, 2, 4])
+
+        self.assertEqual([], subgraph.weighted_edge_list())  # No edges between selected nodes
+        self.assertEqual(["a", "c", "e"], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 0, 1: 2, 2: 4})
+
+    def test_subgraph_with_nodemap_edge_cases(self):
+        graph = rustworkx.PyGraph()
+        graph.add_nodes_from(["a", "b", "c"])
+        graph.add_edges_from([(0, 1, 1), (1, 2, 2)])
+
+        # Test empty node list
+        subgraph, node_map = graph.subgraph_with_nodemap([])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(0, len(subgraph))
+        self.assertEqual(dict(node_map), {})
+
+        # Test invalid node indices (should be silently ignored)
+        subgraph, node_map = graph.subgraph_with_nodemap([42, 100])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(0, len(subgraph))
+        self.assertEqual(dict(node_map), {})
+
+        # Test single node (no edges in subgraph)
+        subgraph, node_map = graph.subgraph_with_nodemap([1])
+        self.assertEqual([], subgraph.weighted_edge_list())
+        self.assertEqual(["b"], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 1})
+
+        # Test all nodes
+        subgraph, node_map = graph.subgraph_with_nodemap([0, 1, 2])
+        self.assertEqual([(0, 1, 1), (1, 2, 2)], subgraph.weighted_edge_list())
+        self.assertEqual(["a", "b", "c"], subgraph.nodes())
+        self.assertEqual(dict(node_map), {0: 0, 1: 1, 2: 2})