annotate Graph.py @ 31:a2358c64d9af

Chapter 5.4, exercise 5: Pareto distribution and city populations.
author Brian Neal <bgneal@gmail.com>
date Mon, 07 Jan 2013 20:41:44 -0600
parents 84e183b40c63
children 10db8c3a6b83
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@8 10 from collections import deque
bgneal@1 11
bgneal@1 12
bgneal@5 13 class GraphError(Exception):
bgneal@5 14 """Exception for Graph errors"""
bgneal@5 15 pass
bgneal@5 16
bgneal@5 17
bgneal@1 18 class Vertex(object):
bgneal@1 19 """A Vertex is a node in a graph."""
bgneal@1 20
bgneal@1 21 def __init__(self, label=''):
bgneal@1 22 self.label = label
bgneal@1 23
bgneal@1 24 def __repr__(self):
bgneal@1 25 """Returns a string representation of this object that can
bgneal@1 26 be evaluated as a Python expression."""
bgneal@1 27 return 'Vertex(%s)' % repr(self.label)
bgneal@1 28
bgneal@1 29 __str__ = __repr__
bgneal@1 30 """The str and repr forms of this object are the same."""
bgneal@1 31
bgneal@1 32
bgneal@1 33 class Edge(tuple):
bgneal@1 34 """An Edge is a list of two vertices."""
bgneal@1 35
bgneal@1 36 def __new__(cls, *vs):
bgneal@1 37 """The Edge constructor takes two vertices."""
bgneal@1 38 if len(vs) != 2:
bgneal@1 39 raise ValueError, 'Edges must connect exactly two vertices.'
bgneal@1 40 return tuple.__new__(cls, vs)
bgneal@1 41
bgneal@1 42 def __repr__(self):
bgneal@1 43 """Return a string representation of this object that can
bgneal@1 44 be evaluated as a Python expression."""
bgneal@1 45 return 'Edge(%s, %s)' % (repr(self[0]), repr(self[1]))
bgneal@1 46
bgneal@1 47 __str__ = __repr__
bgneal@1 48 """The str and repr forms of this object are the same."""
bgneal@1 49
bgneal@1 50
bgneal@1 51 class Graph(dict):
bgneal@1 52 """A Graph is a dictionary of dictionaries. The outer
bgneal@1 53 dictionary maps from a vertex to an inner dictionary.
bgneal@1 54 The inner dictionary maps from other vertices to edges.
bgneal@1 55
bgneal@1 56 For vertices a and b, graph[a][b] maps
bgneal@1 57 to the edge that connects a->b, if it exists."""
bgneal@1 58
bgneal@6 59 def __init__(self, vs=None, es=None):
bgneal@1 60 """Creates a new graph.
bgneal@1 61 vs: list of vertices;
bgneal@1 62 es: list of edges.
bgneal@1 63 """
bgneal@6 64 if vs:
bgneal@6 65 for v in vs:
bgneal@6 66 self.add_vertex(v)
bgneal@1 67
bgneal@6 68 if es:
bgneal@6 69 for e in es:
bgneal@6 70 self.add_edge(e)
bgneal@1 71
bgneal@1 72 def add_vertex(self, v):
bgneal@1 73 """Add a vertex to the graph."""
bgneal@1 74 self[v] = {}
bgneal@1 75
bgneal@1 76 def add_edge(self, e):
bgneal@1 77 """Adds and edge to the graph by adding an entry in both directions.
bgneal@1 78
bgneal@1 79 If there is already an edge connecting these Vertices, the
bgneal@1 80 new edge replaces it.
bgneal@1 81 """
bgneal@1 82 v, w = e
bgneal@1 83 self[v][w] = e
bgneal@1 84 self[w][v] = e
bgneal@1 85
bgneal@1 86 def get_edge(self, v, w):
bgneal@1 87 """Returns the edge object that exists between the two vertices v & w,
bgneal@1 88 or None if no such edge exists.
bgneal@1 89
bgneal@1 90 """
bgneal@1 91 try:
bgneal@1 92 return self[v][w]
bgneal@1 93 except KeyError:
bgneal@1 94 return None
bgneal@1 95
bgneal@1 96 def remove_edge(self, e):
bgneal@1 97 """Removes the edge e from the graph."""
bgneal@1 98
bgneal@4 99 v, w = e
bgneal@4 100 del self[v][w]
bgneal@4 101 del self[w][v]
bgneal@1 102
bgneal@1 103 def vertices(self):
bgneal@1 104 """Returns a list of the vertices in the graph."""
bgneal@1 105
bgneal@1 106 return self.keys()
bgneal@1 107
bgneal@1 108 def edges(self):
bgneal@1 109 """"Returns a list of the edges in the graph."""
bgneal@1 110
bgneal@1 111 edge_set = set()
bgneal@4 112 for d in self.itervalues():
bgneal@4 113 edge_set.update(d.itervalues())
bgneal@1 114
bgneal@1 115 return list(edge_set)
bgneal@1 116
bgneal@1 117 def out_vertices(self, v):
bgneal@1 118 """Returns a list of vertices that are adjacent to the given vertex v.
bgneal@1 119
bgneal@1 120 """
bgneal@1 121 return self[v].keys()
bgneal@1 122
bgneal@1 123 def out_edges(self, v):
bgneal@1 124 """Returns a list of edges connected to a given vertex v."""
bgneal@1 125
bgneal@1 126 return self[v].values()
bgneal@1 127
bgneal@5 128 def remove_all_edges(self):
bgneal@5 129 """Removes all edges in the graph."""
bgneal@5 130 for v in self.iterkeys():
bgneal@5 131 self[v] = {}
bgneal@5 132
bgneal@1 133 def add_all_edges(self):
bgneal@1 134 """Makes the graph complete by adding edges between all pairs of
bgneal@1 135 vertices.
bgneal@1 136
bgneal@1 137 """
bgneal@1 138 # Clear all edges first
bgneal@5 139 self.remove_all_edges()
bgneal@1 140
bgneal@5 141 # For each combination of 2 vertices, create an edge between them:
bgneal@1 142 for v, w in itertools.combinations(self.iterkeys(), 2):
bgneal@6 143 self.add_edge(Edge(v, w))
bgneal@1 144
bgneal@5 145 def add_regular_edges(self, k):
bgneal@5 146 """Makes the graph regular by making every vertex have k edges.
bgneal@5 147
bgneal@5 148 It is not always possible to create a regular graph with a given degree.
bgneal@5 149 If a graph has n vertices, then a regular graph can be constructed with
bgneal@5 150 degree k if n >= k + 1 and n * k is even. If these conditions are not
bgneal@5 151 met a GraphError exception is raised.
bgneal@5 152
bgneal@5 153 """
bgneal@5 154 n = len(self.vertices())
bgneal@5 155 if n < k + 1:
bgneal@5 156 raise GraphError("Can't make a regular graph with degree >= number"
bgneal@5 157 " of vertices")
bgneal@5 158 if (n * k) % 2 != 0:
bgneal@5 159 raise GraphError("Can't make a regular graph of degree k and"
bgneal@5 160 " order n where k * n is odd")
bgneal@5 161
bgneal@5 162 # Remove all edges first
bgneal@5 163 self.remove_all_edges()
bgneal@5 164
bgneal@5 165 if k % 2 != 0: # if k is odd
bgneal@5 166 self._add_regular_edges_even(k - 1)
bgneal@5 167 self._add_regular_edges_odd()
bgneal@5 168 else:
bgneal@5 169 self._add_regular_edges_even(k)
bgneal@5 170
bgneal@5 171 def _add_regular_edges_even(self, k):
bgneal@5 172 """Make a regular graph with degree k. k must be even."""
bgneal@5 173
bgneal@5 174 vs = self.vertices()
bgneal@5 175 vs2 = vs * 2
bgneal@5 176
bgneal@5 177 for i, v in enumerate(vs):
bgneal@5 178 for j in range(1, k / 2 + 1):
bgneal@5 179 w = vs2[i + j]
bgneal@5 180 self.add_edge(Edge(v, w))
bgneal@5 181
bgneal@5 182 def _add_regular_edges_odd(self):
bgneal@5 183 """Adds an extra edge across the graph to finish off a regular graph
bgneal@5 184 with odd degree. The number of vertices must be even.
bgneal@5 185
bgneal@5 186 """
bgneal@5 187 vs = self.vertices()
bgneal@5 188 vs2 = vs * 2
bgneal@5 189 n = len(vs)
bgneal@5 190
bgneal@5 191 for i in range(n / 2):
bgneal@5 192 v = vs2[i]
bgneal@5 193 w = vs2[i + n / 2]
bgneal@5 194 self.add_edge(Edge(v, w))
bgneal@5 195
bgneal@7 196 def bfs(self, start, visit_func=None):
bgneal@7 197 """Perform a breadth first search starting at node start.
bgneal@7 198
bgneal@7 199 The function visit_func, if supplied, is invoked on each node.
bgneal@7 200
bgneal@8 201 The set of visited nodes is returned.
bgneal@8 202
bgneal@7 203 """
bgneal@8 204 visited = set()
bgneal@7 205
bgneal@7 206 # Create a work queue consisting initially of the starting node
bgneal@8 207 queue = deque([start])
bgneal@7 208
bgneal@7 209 while queue:
bgneal@7 210 # retrieve first item from the queue
bgneal@8 211 v = queue.popleft()
bgneal@7 212
bgneal@8 213 if v in visited:
bgneal@7 214 continue # Skip this one if we've seen it before
bgneal@7 215
bgneal@7 216 # Mark it as visited and invoke user's function on it
bgneal@8 217 visited.add(v)
bgneal@7 218 if visit_func:
bgneal@7 219 visit_func(v)
bgneal@7 220
bgneal@7 221 # Add the adjacent neigbors to the node to the queue
bgneal@22 222 queue.extend(c for c in self.out_vertices(v) if c not in visited)
bgneal@7 223
bgneal@8 224 return visited
bgneal@8 225
bgneal@7 226 def is_connected(self):
bgneal@7 227 """Returns True if the graph is connected (there is a path from every
bgneal@7 228 node to every other node) and False otherwise.
bgneal@7 229
bgneal@7 230 """
bgneal@7 231 vs = self.vertices()
bgneal@7 232 if len(vs):
bgneal@8 233 visited = self.bfs(vs[0])
bgneal@7 234 # See if all nodes have been visited
bgneal@8 235 return len(vs) == len(visited)
bgneal@7 236
bgneal@7 237 return False # Graph is empty
bgneal@7 238
bgneal@1 239
bgneal@1 240 def main(script, *args):
bgneal@1 241 import pprint
bgneal@1 242
bgneal@1 243 v = Vertex('v')
bgneal@1 244 print v
bgneal@1 245 w = Vertex('w')
bgneal@1 246 print w
bgneal@1 247 e = Edge(v, w)
bgneal@1 248 print e
bgneal@1 249 g = Graph([v,w], [e])
bgneal@1 250 pprint.pprint(g)
bgneal@1 251
bgneal@1 252 print "g.get_edge(v, w): ", g.get_edge(v, w)
bgneal@1 253 x = Vertex('x')
bgneal@1 254 print "g.get_edge(v, x): ", g.get_edge(v, x)
bgneal@1 255
bgneal@1 256 g.remove_edge(e)
bgneal@1 257 pprint.pprint(g)
bgneal@1 258
bgneal@1 259 print "vertices: ", g.vertices()
bgneal@1 260 print "edges: ", g.edges()
bgneal@1 261
bgneal@1 262 g.add_edge(e)
bgneal@1 263 u = Vertex('u')
bgneal@1 264 e1 = Edge(u, v)
bgneal@1 265 e2 = Edge(u, w)
bgneal@1 266 g.add_vertex(u)
bgneal@1 267 g.add_edge(e1)
bgneal@1 268 g.add_edge(e2)
bgneal@1 269 print "Adding vertex u and edges:"
bgneal@1 270 pprint.pprint(g)
bgneal@1 271 print "vertices: ", g.vertices()
bgneal@1 272 print "edges: ", g.edges()
bgneal@1 273
bgneal@1 274 print "Out vertices for v: ", g.out_vertices(v)
bgneal@1 275 print "Out edges for v: ", g.out_edges(v)
bgneal@1 276
bgneal@1 277 x = Vertex('x')
bgneal@1 278 g.add_vertex(x)
bgneal@1 279 g.add_all_edges()
bgneal@1 280 pprint.pprint(g)
bgneal@1 281
bgneal@7 282 print "g is connected?", g.is_connected()
bgneal@7 283 edges = g.out_edges(v)
bgneal@7 284 for e in edges:
bgneal@7 285 g.remove_edge(e)
bgneal@7 286 pprint.pprint(g)
bgneal@7 287 print "g is connected?", g.is_connected()
bgneal@7 288
bgneal@7 289 # Isolate v and check is_connected() again
bgneal@7 290
bgneal@1 291
bgneal@1 292 if __name__ == '__main__':
bgneal@1 293 import sys
bgneal@1 294 main(*sys.argv)