bgneal@1: """ Code example from Complexity and Computation, a book about
bgneal@1: exploring complexity science with Python.  Available free from
bgneal@1: 
bgneal@1: http://greenteapress.com/complexity
bgneal@1: 
bgneal@1: Copyright 2011 Allen B. Downey.
bgneal@1: Distributed under the GNU General Public License at gnu.org/licenses/gpl.html.
bgneal@1: """
bgneal@1: import itertools
bgneal@1: 
bgneal@1: 
bgneal@5: class GraphError(Exception):
bgneal@5:     """Exception for Graph errors"""
bgneal@5:     pass
bgneal@5: 
bgneal@5: 
bgneal@1: class Vertex(object):
bgneal@1:     """A Vertex is a node in a graph."""
bgneal@1: 
bgneal@1:     def __init__(self, label=''):
bgneal@1:         self.label = label
bgneal@1: 
bgneal@1:     def __repr__(self):
bgneal@1:         """Returns a string representation of this object that can
bgneal@1:         be evaluated as a Python expression."""
bgneal@1:         return 'Vertex(%s)' % repr(self.label)
bgneal@1: 
bgneal@1:     __str__ = __repr__
bgneal@1:     """The str and repr forms of this object are the same."""
bgneal@1: 
bgneal@1: 
bgneal@1: class Edge(tuple):
bgneal@1:     """An Edge is a list of two vertices."""
bgneal@1: 
bgneal@1:     def __new__(cls, *vs):
bgneal@1:         """The Edge constructor takes two vertices."""
bgneal@1:         if len(vs) != 2:
bgneal@1:             raise ValueError, 'Edges must connect exactly two vertices.'
bgneal@1:         return tuple.__new__(cls, vs)
bgneal@1: 
bgneal@1:     def __repr__(self):
bgneal@1:         """Return a string representation of this object that can
bgneal@1:         be evaluated as a Python expression."""
bgneal@1:         return 'Edge(%s, %s)' % (repr(self[0]), repr(self[1]))
bgneal@1: 
bgneal@1:     __str__ = __repr__
bgneal@1:     """The str and repr forms of this object are the same."""
bgneal@1: 
bgneal@1: 
bgneal@1: class Graph(dict):
bgneal@1:     """A Graph is a dictionary of dictionaries.  The outer
bgneal@1:     dictionary maps from a vertex to an inner dictionary.
bgneal@1:     The inner dictionary maps from other vertices to edges.
bgneal@1: 
bgneal@1:     For vertices a and b, graph[a][b] maps
bgneal@1:     to the edge that connects a->b, if it exists."""
bgneal@1: 
bgneal@6:     def __init__(self, vs=None, es=None):
bgneal@1:         """Creates a new graph.
bgneal@1:         vs: list of vertices;
bgneal@1:         es: list of edges.
bgneal@1:         """
bgneal@6:         if vs:
bgneal@6:             for v in vs:
bgneal@6:                 self.add_vertex(v)
bgneal@1: 
bgneal@6:         if es:
bgneal@6:             for e in es:
bgneal@6:                 self.add_edge(e)
bgneal@1: 
bgneal@1:     def add_vertex(self, v):
bgneal@1:         """Add a vertex to the graph."""
bgneal@1:         self[v] = {}
bgneal@1: 
bgneal@1:     def add_edge(self, e):
bgneal@1:         """Adds and edge to the graph by adding an entry in both directions.
bgneal@1: 
bgneal@1:         If there is already an edge connecting these Vertices, the
bgneal@1:         new edge replaces it.
bgneal@1:         """
bgneal@1:         v, w = e
bgneal@1:         self[v][w] = e
bgneal@1:         self[w][v] = e
bgneal@1: 
bgneal@1:     def get_edge(self, v, w):
bgneal@1:         """Returns the edge object that exists between the two vertices v & w,
bgneal@1:         or None if no such edge exists.
bgneal@1: 
bgneal@1:         """
bgneal@1:         try:
bgneal@1:             return self[v][w]
bgneal@1:         except KeyError:
bgneal@1:             return None
bgneal@1: 
bgneal@1:     def remove_edge(self, e):
bgneal@1:         """Removes the edge e from the graph."""
bgneal@1: 
bgneal@4:         v, w = e
bgneal@4:         del self[v][w]
bgneal@4:         del self[w][v]
bgneal@1: 
bgneal@1:     def vertices(self):
bgneal@1:         """Returns a list of the vertices in the graph."""
bgneal@1: 
bgneal@1:         return self.keys()
bgneal@1: 
bgneal@1:     def edges(self):
bgneal@1:         """"Returns a list of the edges in the graph."""
bgneal@1: 
bgneal@1:         edge_set = set()
bgneal@4:         for d in self.itervalues():
bgneal@4:             edge_set.update(d.itervalues())
bgneal@1: 
bgneal@1:         return list(edge_set)
bgneal@1: 
bgneal@1:     def out_vertices(self, v):
bgneal@1:         """Returns a list of vertices that are adjacent to the given vertex v.
bgneal@1: 
bgneal@1:         """
bgneal@1:         return self[v].keys()
bgneal@1: 
bgneal@1:     def out_edges(self, v):
bgneal@1:         """Returns a list of edges connected to a given vertex v."""
bgneal@1: 
bgneal@1:         return self[v].values()
bgneal@1: 
bgneal@5:     def remove_all_edges(self):
bgneal@5:         """Removes all edges in the graph."""
bgneal@5:         for v in self.iterkeys():
bgneal@5:             self[v] = {}
bgneal@5: 
bgneal@1:     def add_all_edges(self):
bgneal@1:         """Makes the graph complete by adding edges between all pairs of
bgneal@1:         vertices.
bgneal@1: 
bgneal@1:         """
bgneal@1:         # Clear all edges first
bgneal@5:         self.remove_all_edges()
bgneal@1: 
bgneal@5:         # For each combination of 2 vertices, create an edge between them:
bgneal@1:         for v, w in itertools.combinations(self.iterkeys(), 2):
bgneal@6:             self.add_edge(Edge(v, w))
bgneal@1: 
bgneal@5:     def add_regular_edges(self, k):
bgneal@5:         """Makes the graph regular by making every vertex have k edges.
bgneal@5: 
bgneal@5:         It is not always possible to create a regular graph with a given degree.
bgneal@5:         If a graph has n vertices, then a regular graph can be constructed with
bgneal@5:         degree k if n >= k + 1 and n * k is even. If these conditions are not
bgneal@5:         met a GraphError exception is raised.
bgneal@5: 
bgneal@5:         """
bgneal@5:         n = len(self.vertices())
bgneal@5:         if n < k + 1:
bgneal@5:             raise GraphError("Can't make a regular graph with degree >= number"
bgneal@5:                              " of vertices")
bgneal@5:         if (n * k) % 2 != 0:
bgneal@5:             raise GraphError("Can't make a regular graph of degree k and"
bgneal@5:                              " order n where k * n is odd")
bgneal@5: 
bgneal@5:         # Remove all edges first
bgneal@5:         self.remove_all_edges()
bgneal@5: 
bgneal@5:         if k % 2 != 0:      # if k is odd
bgneal@5:             self._add_regular_edges_even(k - 1)
bgneal@5:             self._add_regular_edges_odd()
bgneal@5:         else:
bgneal@5:             self._add_regular_edges_even(k)
bgneal@5: 
bgneal@5:     def _add_regular_edges_even(self, k):
bgneal@5:         """Make a regular graph with degree k. k must be even."""
bgneal@5: 
bgneal@5:         vs = self.vertices()
bgneal@5:         vs2 = vs * 2
bgneal@5: 
bgneal@5:         for i, v in enumerate(vs):
bgneal@5:             for j in range(1, k / 2 + 1):
bgneal@5:                 w = vs2[i + j]
bgneal@5:                 self.add_edge(Edge(v, w))
bgneal@5: 
bgneal@5:     def _add_regular_edges_odd(self):
bgneal@5:         """Adds an extra edge across the graph to finish off a regular graph
bgneal@5:         with odd degree. The number of vertices must be even.
bgneal@5: 
bgneal@5:         """
bgneal@5:         vs = self.vertices()
bgneal@5:         vs2 = vs * 2
bgneal@5:         n = len(vs)
bgneal@5: 
bgneal@5:         for i in range(n / 2):
bgneal@5:             v = vs2[i]
bgneal@5:             w = vs2[i + n / 2]
bgneal@5:             self.add_edge(Edge(v, w))
bgneal@5: 
bgneal@7:     def bfs(self, start, visit_func=None):
bgneal@7:         """Perform a breadth first search starting at node start.
bgneal@7: 
bgneal@7:         The function visit_func, if supplied, is invoked on each node.
bgneal@7: 
bgneal@7:         """
bgneal@7:         # Mark all vertices as unvisited
bgneal@7:         vs = self.vertices()
bgneal@7:         for v in vs:
bgneal@7:             v.visited = False
bgneal@7: 
bgneal@7:         # Create a work queue consisting initially of the starting node
bgneal@7:         queue = [start]
bgneal@7: 
bgneal@7:         while queue:
bgneal@7:             # retrieve first item from the queue
bgneal@7:             v = queue.pop(0)
bgneal@7: 
bgneal@7:             if v.visited:
bgneal@7:                 continue            # Skip this one if we've seen it before
bgneal@7: 
bgneal@7:             # Mark it as visited and invoke user's function on it
bgneal@7:             v.visited = True
bgneal@7:             if visit_func:
bgneal@7:                 visit_func(v)
bgneal@7: 
bgneal@7:             # Add the adjacent neigbors to the node to the queue
bgneal@7:             queue.extend(self.out_vertices(v))
bgneal@7: 
bgneal@7:     def is_connected(self):
bgneal@7:         """Returns True if the graph is connected (there is a path from every
bgneal@7:         node to every other node) and False otherwise.
bgneal@7: 
bgneal@7:         """
bgneal@7:         vs = self.vertices()
bgneal@7:         if len(vs):
bgneal@7:             self.bfs(vs[0])
bgneal@7:             # See if all nodes have been visited
bgneal@7:             for v in vs:
bgneal@7:                 if not v.visited:
bgneal@7:                     return False
bgneal@7:             return True
bgneal@7: 
bgneal@7:         return False        # Graph is empty
bgneal@7: 
bgneal@1: 
bgneal@1: def main(script, *args):
bgneal@1:     import pprint
bgneal@1: 
bgneal@1:     v = Vertex('v')
bgneal@1:     print v
bgneal@1:     w = Vertex('w')
bgneal@1:     print w
bgneal@1:     e = Edge(v, w)
bgneal@1:     print e
bgneal@1:     g = Graph([v,w], [e])
bgneal@1:     pprint.pprint(g)
bgneal@1: 
bgneal@1:     print "g.get_edge(v, w): ", g.get_edge(v, w)
bgneal@1:     x = Vertex('x')
bgneal@1:     print "g.get_edge(v, x): ", g.get_edge(v, x)
bgneal@1: 
bgneal@1:     g.remove_edge(e)
bgneal@1:     pprint.pprint(g)
bgneal@1: 
bgneal@1:     print "vertices: ", g.vertices()
bgneal@1:     print "edges: ", g.edges()
bgneal@1: 
bgneal@1:     g.add_edge(e)
bgneal@1:     u = Vertex('u')
bgneal@1:     e1 = Edge(u, v)
bgneal@1:     e2 = Edge(u, w)
bgneal@1:     g.add_vertex(u)
bgneal@1:     g.add_edge(e1)
bgneal@1:     g.add_edge(e2)
bgneal@1:     print "Adding vertex u and edges:"
bgneal@1:     pprint.pprint(g)
bgneal@1:     print "vertices: ", g.vertices()
bgneal@1:     print "edges: ", g.edges()
bgneal@1: 
bgneal@1:     print "Out vertices for v: ", g.out_vertices(v)
bgneal@1:     print "Out edges for v: ", g.out_edges(v)
bgneal@1: 
bgneal@1:     x = Vertex('x')
bgneal@1:     g.add_vertex(x)
bgneal@1:     g.add_all_edges()
bgneal@1:     pprint.pprint(g)
bgneal@1: 
bgneal@7:     print "g is connected?", g.is_connected()
bgneal@7:     edges = g.out_edges(v)
bgneal@7:     for e in edges:
bgneal@7:         g.remove_edge(e)
bgneal@7:     pprint.pprint(g)
bgneal@7:     print "g is connected?", g.is_connected()
bgneal@7: 
bgneal@7:     # Isolate v and check is_connected() again
bgneal@7: 
bgneal@1: 
bgneal@1: if __name__ == '__main__':
bgneal@1:     import sys
bgneal@1:     main(*sys.argv)