diff Graph.py @ 8:1defe6fcf9d3

Improvements to Graph's bfs.
author Brian Neal <bgneal@gmail.com>
date Mon, 03 Dec 2012 19:11:17 -0600
parents 69e5977417b3
children 84e183b40c63
line wrap: on
line diff
--- a/Graph.py	Sun Dec 02 21:51:51 2012 -0600
+++ b/Graph.py	Mon Dec 03 19:11:17 2012 -0600
@@ -7,6 +7,7 @@
 Distributed under the GNU General Public License at gnu.org/licenses/gpl.html.
 """
 import itertools
+from collections import deque
 
 
 class GraphError(Exception):
@@ -197,30 +198,31 @@
 
         The function visit_func, if supplied, is invoked on each node.
 
+        The set of visited nodes is returned.
+
         """
-        # Mark all vertices as unvisited
-        vs = self.vertices()
-        for v in vs:
-            v.visited = False
+        visited = set()
 
         # Create a work queue consisting initially of the starting node
-        queue = [start]
+        queue = deque([start])
 
         while queue:
             # retrieve first item from the queue
-            v = queue.pop(0)
+            v = queue.popleft()
 
-            if v.visited:
+            if v in visited:
                 continue            # Skip this one if we've seen it before
 
             # Mark it as visited and invoke user's function on it
-            v.visited = True
+            visited.add(v)
             if visit_func:
                 visit_func(v)
 
             # Add the adjacent neigbors to the node to the queue
             queue.extend(self.out_vertices(v))
 
+        return visited
+
     def is_connected(self):
         """Returns True if the graph is connected (there is a path from every
         node to every other node) and False otherwise.
@@ -228,12 +230,9 @@
         """
         vs = self.vertices()
         if len(vs):
-            self.bfs(vs[0])
+            visited = self.bfs(vs[0])
             # See if all nodes have been visited
-            for v in vs:
-                if not v.visited:
-                    return False
-            return True
+            return len(vs) == len(visited)
 
         return False        # Graph is empty