Add neighbors methods to graph classes

src/digraph.rs

@@ -1076,6 +1076,61 @@ impl PyDiGraph {
         Ok(out_map)
     }
 
+    /// Get the neighbors of a node.
+    ///
+    /// This will return a list of neighbor node indices.
+    ///
+    /// :param int node: The index of the node to get the neighbors of
+    ///
+    /// :returns: A list of the neighbor node indicies
+    /// :rtype: list
+    #[text_signature = "(node, /)"]
+    pub fn neighbors(&self, node: usize) -> Vec<usize> {
+        self.graph
+            .neighbors(NodeIndex::new(node))
+            .map(|node| node.index())
+            .collect()
+    }
+
+    /// Get the successor indices of a node.
+    ///
+    /// This will return a list of the node indicies for the succesors of
+    /// a node
+    ///
+    /// :param int node: The index of the node to get the successors of
+    ///
+    /// :returns: A list of the neighbor node indicies
+    /// :rtype: list
+    #[text_signature = "(node, /)"]
+    pub fn successor_indices(&mut self, node: usize) -> Vec<usize> {
+        self.graph
+            .neighbors_directed(
+                NodeIndex::new(node),
+                petgraph::Direction::Outgoing,
+            )
+            .map(|node| node.index())
+            .collect()
+    }
+
+    /// Get the predecessor indices of a node.
+    ///
+    /// This will return a list of the node indicies for the predecessors of
+    /// a node
+    ///
+    /// :param int node: The index of the node to get the predecessors of
+    ///
+    /// :returns: A list of the neighbor node indicies
+    /// :rtype: list
+    #[text_signature = "(node, /)"]
+    pub fn predecessor_indices(&mut self, node: usize) -> Vec<usize> {
+        self.graph
+            .neighbors_directed(
+                NodeIndex::new(node),
+                petgraph::Direction::Incoming,
+            )
+            .map(|node| node.index())
+            .collect()
+    }
     /// Get the index and edge data for all parents of a node.
     ///
     /// This will return a list of tuples with the parent index the node index

src/graph.rs

@@ -777,6 +777,22 @@ impl PyGraph {
         Ok(out_map)
     }
 
+    /// Get the neighbors of a node.
+    ///
+    /// This with return a list of neighbor node indices
+    ///
+    /// :param int node: The index of the node to get the neibhors of
+    ///
+    /// :returns: A list of the neighbor node indicies
+    /// :rtype: list
+    #[text_signature = "(node, /)"]
+    pub fn neighbors(&self, node: usize) -> Vec<usize> {
+        self.graph
+            .neighbors(NodeIndex::new(node))
+            .map(|node| node.index())
+            .collect()
+    }
+
     /// Get the degree for a node
     ///
     /// :param int node: The index of the  node to find the inbound degree of

tests/graph/test_neighbors.py

@@ -0,0 +1,32 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import retworkx
+
+
+class TestNeighbors(unittest.TestCase):
+    def test_single_neighbor(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        node_b = graph.add_node('b')
+        graph.add_edge(node_a, node_b, {'a': 1})
+        node_c = graph.add_node('c')
+        graph.add_edge(node_a, node_c, {'a': 2})
+        res = graph.neighbors(node_a)
+        self.assertEqual([node_c, node_b], res)
+
+    def test_no_neighbor(self):
+        graph = retworkx.PyGraph()
+        node_a = graph.add_node('a')
+        self.assertEqual([], graph.neighbors(node_a))

tests/test_neighbors.py

@@ -0,0 +1,50 @@
+# Licensed under the Apache License, Version 2.0 (the "License"); you may
+# not use this file except in compliance with the License. You may obtain
+# a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+# License for the specific language governing permissions and limitations
+# under the License.
+
+import unittest
+
+import retworkx
+
+
+class TestAdj(unittest.TestCase):
+    def test_single_neighbor(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        node_b = dag.add_child(node_a, 'b', {'a': 1})
+        node_c = dag.add_child(node_a, 'c', {'a': 2})
+        res = dag.neighbors(node_a)
+        self.assertEqual([node_c, node_b], res)
+
+    def test_single_neighbor_dir(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        node_b = dag.add_child(node_a, 'b', {'a': 1})
+        node_c = dag.add_child(node_a, 'c', {'a': 2})
+        res = dag.successor_indices(node_a)
+        self.assertEqual([node_c, node_b], res)
+        res = dag.predecessor_indices(node_a)
+        self.assertEqual([], res)
+
+    def test_neighbor_dir_surrounded(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        node_b = dag.add_child(node_a, 'b', {'a': 1})
+        node_c = dag.add_child(node_b, 'c', {'a': 2})
+        res = dag.successor_indices(node_b)
+        self.assertEqual([node_c], res)
+        res = dag.predecessor_indices(node_b)
+        self.assertEqual([node_a], res)
+
+    def test_no_neighbor(self):
+        dag = retworkx.PyDAG()
+        node_a = dag.add_node('a')
+        self.assertEqual([], dag.neighbors(node_a))