Mercurial > public > think_complexity

annotate Graph.py @ 1:de3ae15ebbfa

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Exercise 2: created Graph.py.
author Brian Neal <bgneal@gmail.com>
date Thu, 29 Nov 2012 20:06:54 -0600
parents
children 9d0cf96b6a3b
rev   line source
bgneal@1 1 """ Code example from Complexity and Computation, a book about
bgneal@1 2 exploring complexity science with Python. Available free from
bgneal@1 3
bgneal@1 4 http://greenteapress.com/complexity
bgneal@1 5
bgneal@1 6 Copyright 2011 Allen B. Downey.
bgneal@1 7 Distributed under the GNU General Public License at gnu.org/licenses/gpl.html.
bgneal@1 8 """
bgneal@1 9 import itertools
bgneal@1 10
bgneal@1 11
bgneal@1 12 class Vertex(object):
bgneal@1 13 """A Vertex is a node in a graph."""
bgneal@1 14
bgneal@1 15 def __init__(self, label=''):
bgneal@1 16 self.label = label
bgneal@1 17
bgneal@1 18 def __repr__(self):
bgneal@1 19 """Returns a string representation of this object that can
bgneal@1 20 be evaluated as a Python expression."""
bgneal@1 21 return 'Vertex(%s)' % repr(self.label)
bgneal@1 22
bgneal@1 23 __str__ = __repr__
bgneal@1 24 """The str and repr forms of this object are the same."""
bgneal@1 25
bgneal@1 26
bgneal@1 27 class Edge(tuple):
bgneal@1 28 """An Edge is a list of two vertices."""
bgneal@1 29
bgneal@1 30 def __new__(cls, *vs):
bgneal@1 31 """The Edge constructor takes two vertices."""
bgneal@1 32 if len(vs) != 2:
bgneal@1 33 raise ValueError, 'Edges must connect exactly two vertices.'
bgneal@1 34 return tuple.__new__(cls, vs)
bgneal@1 35
bgneal@1 36 def __repr__(self):
bgneal@1 37 """Return a string representation of this object that can
bgneal@1 38 be evaluated as a Python expression."""
bgneal@1 39 return 'Edge(%s, %s)' % (repr(self[0]), repr(self[1]))
bgneal@1 40
bgneal@1 41 __str__ = __repr__
bgneal@1 42 """The str and repr forms of this object are the same."""
bgneal@1 43
bgneal@1 44
bgneal@1 45 class Graph(dict):
bgneal@1 46 """A Graph is a dictionary of dictionaries. The outer
bgneal@1 47 dictionary maps from a vertex to an inner dictionary.
bgneal@1 48 The inner dictionary maps from other vertices to edges.
bgneal@1 49
bgneal@1 50 For vertices a and b, graph[a][b] maps
bgneal@1 51 to the edge that connects a->b, if it exists."""
bgneal@1 52
bgneal@1 53 def __init__(self, vs=[], es=[]):
bgneal@1 54 """Creates a new graph.
bgneal@1 55 vs: list of vertices;
bgneal@1 56 es: list of edges.
bgneal@1 57 """
bgneal@1 58 for v in vs:
bgneal@1 59 self.add_vertex(v)
bgneal@1 60
bgneal@1 61 for e in es:
bgneal@1 62 self.add_edge(e)
bgneal@1 63
bgneal@1 64 def add_vertex(self, v):
bgneal@1 65 """Add a vertex to the graph."""
bgneal@1 66 self[v] = {}
bgneal@1 67
bgneal@1 68 def add_edge(self, e):
bgneal@1 69 """Adds and edge to the graph by adding an entry in both directions.
bgneal@1 70
bgneal@1 71 If there is already an edge connecting these Vertices, the
bgneal@1 72 new edge replaces it.
bgneal@1 73 """
bgneal@1 74 v, w = e
bgneal@1 75 self[v][w] = e
bgneal@1 76 self[w][v] = e
bgneal@1 77
bgneal@1 78 def get_edge(self, v, w):
bgneal@1 79 """Returns the edge object that exists between the two vertices v & w,
bgneal@1 80 or None if no such edge exists.
bgneal@1 81
bgneal@1 82 """
bgneal@1 83 try:
bgneal@1 84 return self[v][w]
bgneal@1 85 except KeyError:
bgneal@1 86 return None
bgneal@1 87
bgneal@1 88 def remove_edge(self, e):
bgneal@1 89 """Removes the edge e from the graph."""
bgneal@1 90
bgneal@1 91 for x in self.iterkeys():
bgneal@1 92 remove = [k for k, v in self[x].iteritems() if v is e]
bgneal@1 93 for k in remove:
bgneal@1 94 del self[x][k]
bgneal@1 95
bgneal@1 96 def vertices(self):
bgneal@1 97 """Returns a list of the vertices in the graph."""
bgneal@1 98
bgneal@1 99 return self.keys()
bgneal@1 100
bgneal@1 101 def edges(self):
bgneal@1 102 """"Returns a list of the edges in the graph."""
bgneal@1 103
bgneal@1 104 edge_set = set()
bgneal@1 105 for x in self.iterkeys():
bgneal@1 106 for e in self[x].itervalues():
bgneal@1 107 edge_set.add(e)
bgneal@1 108
bgneal@1 109 return list(edge_set)
bgneal@1 110
bgneal@1 111 def out_vertices(self, v):
bgneal@1 112 """Returns a list of vertices that are adjacent to the given vertex v.
bgneal@1 113
bgneal@1 114 """
bgneal@1 115 return self[v].keys()
bgneal@1 116
bgneal@1 117 def out_edges(self, v):
bgneal@1 118 """Returns a list of edges connected to a given vertex v."""
bgneal@1 119
bgneal@1 120 return self[v].values()
bgneal@1 121
bgneal@1 122 def add_all_edges(self):
bgneal@1 123 """Makes the graph complete by adding edges between all pairs of
bgneal@1 124 vertices.
bgneal@1 125
bgneal@1 126 """
bgneal@1 127 # Clear all edges first
bgneal@1 128 for v in self.iterkeys():
bgneal@1 129 self[v] = {}
bgneal@1 130
bgneal@1 131 for v, w in itertools.combinations(self.iterkeys(), 2):
bgneal@1 132 e = Edge(v, w)
bgneal@1 133 self[v][w] = e
bgneal@1 134 self[w][v] = e
bgneal@1 135
bgneal@1 136
bgneal@1 137 def main(script, *args):
bgneal@1 138 import pprint
bgneal@1 139
bgneal@1 140 v = Vertex('v')
bgneal@1 141 print v
bgneal@1 142 w = Vertex('w')
bgneal@1 143 print w
bgneal@1 144 e = Edge(v, w)
bgneal@1 145 print e
bgneal@1 146 g = Graph([v,w], [e])
bgneal@1 147 pprint.pprint(g)
bgneal@1 148
bgneal@1 149 print "g.get_edge(v, w): ", g.get_edge(v, w)
bgneal@1 150 x = Vertex('x')
bgneal@1 151 print "g.get_edge(v, x): ", g.get_edge(v, x)
bgneal@1 152
bgneal@1 153 g.remove_edge(e)
bgneal@1 154 pprint.pprint(g)
bgneal@1 155
bgneal@1 156 print "vertices: ", g.vertices()
bgneal@1 157 print "edges: ", g.edges()
bgneal@1 158
bgneal@1 159 g.add_edge(e)
bgneal@1 160 u = Vertex('u')
bgneal@1 161 e1 = Edge(u, v)
bgneal@1 162 e2 = Edge(u, w)
bgneal@1 163 g.add_vertex(u)
bgneal@1 164 g.add_edge(e1)
bgneal@1 165 g.add_edge(e2)
bgneal@1 166 print "Adding vertex u and edges:"
bgneal@1 167 pprint.pprint(g)
bgneal@1 168 print "vertices: ", g.vertices()
bgneal@1 169 print "edges: ", g.edges()
bgneal@1 170
bgneal@1 171 print "Out vertices for v: ", g.out_vertices(v)
bgneal@1 172 print "Out edges for v: ", g.out_edges(v)
bgneal@1 173
bgneal@1 174 x = Vertex('x')
bgneal@1 175 g.add_vertex(x)
bgneal@1 176 g.add_all_edges()
bgneal@1 177 pprint.pprint(g)
bgneal@1 178
bgneal@1 179
bgneal@1 180 if __name__ == '__main__':
bgneal@1 181 import sys
bgneal@1 182 main(*sys.argv)