diff Graph.py @ 1:de3ae15ebbfa

Exercise 2: created Graph.py.
author Brian Neal <bgneal@gmail.com>
date Thu, 29 Nov 2012 20:06:54 -0600
parents
children 9d0cf96b6a3b
line wrap: on
line diff
--- /dev/null	Thu Jan 01 00:00:00 1970 +0000
+++ b/Graph.py	Thu Nov 29 20:06:54 2012 -0600
@@ -0,0 +1,182 @@
+""" Code example from Complexity and Computation, a book about
+exploring complexity science with Python.  Available free from
+
+http://greenteapress.com/complexity
+
+Copyright 2011 Allen B. Downey.
+Distributed under the GNU General Public License at gnu.org/licenses/gpl.html.
+"""
+import itertools
+
+
+class Vertex(object):
+    """A Vertex is a node in a graph."""
+
+    def __init__(self, label=''):
+        self.label = label
+
+    def __repr__(self):
+        """Returns a string representation of this object that can
+        be evaluated as a Python expression."""
+        return 'Vertex(%s)' % repr(self.label)
+
+    __str__ = __repr__
+    """The str and repr forms of this object are the same."""
+
+
+class Edge(tuple):
+    """An Edge is a list of two vertices."""
+
+    def __new__(cls, *vs):
+        """The Edge constructor takes two vertices."""
+        if len(vs) != 2:
+            raise ValueError, 'Edges must connect exactly two vertices.'
+        return tuple.__new__(cls, vs)
+
+    def __repr__(self):
+        """Return a string representation of this object that can
+        be evaluated as a Python expression."""
+        return 'Edge(%s, %s)' % (repr(self[0]), repr(self[1]))
+
+    __str__ = __repr__
+    """The str and repr forms of this object are the same."""
+
+
+class Graph(dict):
+    """A Graph is a dictionary of dictionaries.  The outer
+    dictionary maps from a vertex to an inner dictionary.
+    The inner dictionary maps from other vertices to edges.
+
+    For vertices a and b, graph[a][b] maps
+    to the edge that connects a->b, if it exists."""
+
+    def __init__(self, vs=[], es=[]):
+        """Creates a new graph.
+        vs: list of vertices;
+        es: list of edges.
+        """
+        for v in vs:
+            self.add_vertex(v)
+
+        for e in es:
+            self.add_edge(e)
+
+    def add_vertex(self, v):
+        """Add a vertex to the graph."""
+        self[v] = {}
+
+    def add_edge(self, e):
+        """Adds and edge to the graph by adding an entry in both directions.
+
+        If there is already an edge connecting these Vertices, the
+        new edge replaces it.
+        """
+        v, w = e
+        self[v][w] = e
+        self[w][v] = e
+
+    def get_edge(self, v, w):
+        """Returns the edge object that exists between the two vertices v & w,
+        or None if no such edge exists.
+
+        """
+        try:
+            return self[v][w]
+        except KeyError:
+            return None
+
+    def remove_edge(self, e):
+        """Removes the edge e from the graph."""
+
+        for x in self.iterkeys():
+            remove = [k for k, v in self[x].iteritems() if v is e]
+            for k in remove:
+                del self[x][k]
+
+    def vertices(self):
+        """Returns a list of the vertices in the graph."""
+
+        return self.keys()
+
+    def edges(self):
+        """"Returns a list of the edges in the graph."""
+
+        edge_set = set()
+        for x in self.iterkeys():
+            for e in self[x].itervalues():
+                edge_set.add(e)
+
+        return list(edge_set)
+
+    def out_vertices(self, v):
+        """Returns a list of vertices that are adjacent to the given vertex v.
+
+        """
+        return self[v].keys()
+
+    def out_edges(self, v):
+        """Returns a list of edges connected to a given vertex v."""
+
+        return self[v].values()
+
+    def add_all_edges(self):
+        """Makes the graph complete by adding edges between all pairs of
+        vertices.
+
+        """
+        # Clear all edges first
+        for v in self.iterkeys():
+            self[v] = {}
+
+        for v, w in itertools.combinations(self.iterkeys(), 2):
+            e = Edge(v, w)
+            self[v][w] = e
+            self[w][v] = e
+
+
+def main(script, *args):
+    import pprint
+
+    v = Vertex('v')
+    print v
+    w = Vertex('w')
+    print w
+    e = Edge(v, w)
+    print e
+    g = Graph([v,w], [e])
+    pprint.pprint(g)
+
+    print "g.get_edge(v, w): ", g.get_edge(v, w)
+    x = Vertex('x')
+    print "g.get_edge(v, x): ", g.get_edge(v, x)
+
+    g.remove_edge(e)
+    pprint.pprint(g)
+
+    print "vertices: ", g.vertices()
+    print "edges: ", g.edges()
+
+    g.add_edge(e)
+    u = Vertex('u')
+    e1 = Edge(u, v)
+    e2 = Edge(u, w)
+    g.add_vertex(u)
+    g.add_edge(e1)
+    g.add_edge(e2)
+    print "Adding vertex u and edges:"
+    pprint.pprint(g)
+    print "vertices: ", g.vertices()
+    print "edges: ", g.edges()
+
+    print "Out vertices for v: ", g.out_vertices(v)
+    print "Out edges for v: ", g.out_edges(v)
+
+    x = Vertex('x')
+    g.add_vertex(x)
+    g.add_all_edges()
+    pprint.pprint(g)
+
+
+if __name__ == '__main__':
+    import sys
+    main(*sys.argv)