Fix non-linear running time in descendants and ancestors methods.

Author: mdickinson

refcycle/i_directed_graph.py

@@ -164,15 +164,16 @@ def descendants(self, start, generations=None):
 
         """
         visited = self.vertex_set()
+        visited.add(start)
         to_visit = deque([(start, 0)])
         while to_visit:
             vertex, depth = to_visit.popleft()
-            visited.add(vertex)
-            if generations is None or depth < generations:
-                to_visit.extend(
-                    (child, depth+1) for child in self.children(vertex)
-                    if child not in visited
-                )
+            if depth == generations:
+                continue
+            for child in self.children(vertex):
+                if child not in visited:
+                    visited.add(child)
+                    to_visit.append((child, depth+1))
         return self.full_subgraph(visited)
 
     def ancestors(self, start, generations=None):
@@ -185,15 +186,16 @@ def ancestors(self, start, generations=None):
 
         """
         visited = self.vertex_set()
+        visited.add(start)
         to_visit = deque([(start, 0)])
         while to_visit:
             vertex, depth = to_visit.popleft()
-            visited.add(vertex)
-            if generations is None or depth < generations:
-                to_visit.extend(
-                    (parent, depth+1) for parent in self.parents(vertex)
-                    if parent not in visited
-                )
+            if depth == generations:
+                continue
+            for parent in self.parents(vertex):
+                if parent not in visited:
+                    visited.add(parent)
+                    to_visit.append((parent, depth+1))
         return self.full_subgraph(visited)
 
     def shortest_path(self, start, end):

refcycle/test/test_directed_graph.py

@@ -199,6 +199,23 @@ def test_more_limited_descendants(self):
             list('12345678'),
         )
 
+    def test_descendants_slow_case(self):
+        # Regression test for #48. A buggy earlier version of the
+        # descendants method had running time exponential in
+        # vertex_count.
+        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,
+        )
+        descendants = graph.descendants(0)
+        self.assertEqual(set(descendants), vertices)
+
     def test_limited_ancestors(self):
         graph = graph_from_string(
             "1 2 3 4 5 6; 1->2 1->3 2->3 2->4 4->3 4->5 5->2 5->6 6->3 6->4")
@@ -220,6 +237,23 @@ def test_limited_ancestors(self):
             ['1', '2', '3', '4', '5', '6'],
         )
 
+    def test_ancestors_slow_case(self):
+        # Regression test for #48. A buggy earlier version of the
+        # ancestors method had running time exponential in
+        # vertex_count.
+        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,
+        )
+        ancestors = graph.ancestors(0)
+        self.assertEqual(set(ancestors), vertices)
+
     def test_length(self):
         self.assertEqual(len(test_graph), 11)