Mercurial > public > think_complexity
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)