Merge pull request #63 from mdickinson/fix/subgraph-from-list

Fix quadratic-time behaviour in DirectedGraph.full_subgraph.

Author: mdickinson

refcycle/directed_graph.py

@@ -102,9 +102,9 @@ def full_subgraph(self, vertices):
         """
         subgraph_vertices = {v for v in vertices}
         subgraph_edges = {edge
-                          for v in vertices
+                          for v in subgraph_vertices
                           for edge in self._out_edges[v]
-                          if self._heads[edge] in vertices}
+                          if self._heads[edge] in subgraph_vertices}
         subgraph_heads = {edge: self._heads[edge]
                           for edge in subgraph_edges}
         subgraph_tails = {edge: self._tails[edge]

refcycle/test/test_directed_graph.py

@@ -291,6 +291,38 @@ def test_full_subgraph(self):
         self.assertCountEqual(vertices, [1, 2, 3, 4, 5])
         self.assertEqual(len(edges), 5)
 
+    def test_full_subgraph_large_from_list(self):
+        # An earlier version of full_subgraph had quadratic-time behaviour.
+        vertex_count = 20000
+        vertices = set(range(vertex_count))
+        edge_mapper = {
+            n: [(n + 1) % vertex_count, (n + 1) % vertex_count]
+            for n in vertices
+        }
+        graph = DirectedGraph.from_out_edges(
+            vertices=vertices,
+            edge_mapper=edge_mapper,
+        )
+        subgraph = graph.full_subgraph(list(vertices))
+        self.assertEqual(len(subgraph.vertices), len(graph.vertices))
+        self.assertEqual(len(subgraph.edges), len(graph.edges))
+
+    def test_full_subgraph_from_iterator(self):
+        # Should be fine to create a subgraph from an iterator.
+        vertex_count = 100
+        vertices = set(range(vertex_count))
+        edge_mapper = {
+            n: [(n + 1) % vertex_count, (n + 1) % vertex_count]
+            for n in vertices
+        }
+        graph = DirectedGraph.from_out_edges(
+            vertices=vertices,
+            edge_mapper=edge_mapper,
+        )
+        subgraph = graph.full_subgraph(iter(vertices))
+        self.assertEqual(len(subgraph.vertices), len(graph.vertices))
+        self.assertEqual(len(subgraph.edges), len(graph.edges))
+
     def test_to_dot(self):
         dot = test_graph.to_dot()
         self.assertIsInstance(dot, six.text_type)