Mercurial > public > think_complexity

annotate Graph.py @ 4:9d0cf96b6a3b

Find changesets by keywords (author, files, the commit message), revision number or hash, or revset expression.
Updates to Graph.py after looking at Prof. Downey's code.
author Brian Neal <bgneal@gmail.com>
date Thu, 29 Nov 2012 20:37:03 -0600
parents de3ae15ebbfa
children 8e44660965ef
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@4 91 v, w = e
bgneal@4 92 del self[v][w]
bgneal@4 93 del self[w][v]
bgneal@1 94
bgneal@1 95 def vertices(self):
bgneal@1 96 """Returns a list of the vertices in the graph."""
bgneal@1 97
bgneal@1 98 return self.keys()
bgneal@1 99
bgneal@1 100 def edges(self):
bgneal@1 101 """"Returns a list of the edges in the graph."""
bgneal@1 102
bgneal@1 103 edge_set = set()
bgneal@4 104 for d in self.itervalues():
bgneal@4 105 edge_set.update(d.itervalues())
bgneal@1 106
bgneal@1 107 return list(edge_set)
bgneal@1 108
bgneal@1 109 def out_vertices(self, v):
bgneal@1 110 """Returns a list of vertices that are adjacent to the given vertex v.
bgneal@1 111
bgneal@1 112 """
bgneal@1 113 return self[v].keys()
bgneal@1 114
bgneal@1 115 def out_edges(self, v):
bgneal@1 116 """Returns a list of edges connected to a given vertex v."""
bgneal@1 117
bgneal@1 118 return self[v].values()
bgneal@1 119
bgneal@1 120 def add_all_edges(self):
bgneal@1 121 """Makes the graph complete by adding edges between all pairs of
bgneal@1 122 vertices.
bgneal@1 123
bgneal@1 124 """
bgneal@1 125 # Clear all edges first
bgneal@1 126 for v in self.iterkeys():
bgneal@1 127 self[v] = {}
bgneal@1 128
bgneal@1 129 for v, w in itertools.combinations(self.iterkeys(), 2):
bgneal@1 130 e = Edge(v, w)
bgneal@1 131 self[v][w] = e
bgneal@1 132 self[w][v] = e
bgneal@1 133
bgneal@1 134
bgneal@1 135 def main(script, *args):
bgneal@1 136 import pprint
bgneal@1 137
bgneal@1 138 v = Vertex('v')
bgneal@1 139 print v
bgneal@1 140 w = Vertex('w')
bgneal@1 141 print w
bgneal@1 142 e = Edge(v, w)
bgneal@1 143 print e
bgneal@1 144 g = Graph([v,w], [e])
bgneal@1 145 pprint.pprint(g)
bgneal@1 146
bgneal@1 147 print "g.get_edge(v, w): ", g.get_edge(v, w)
bgneal@1 148 x = Vertex('x')
bgneal@1 149 print "g.get_edge(v, x): ", g.get_edge(v, x)
bgneal@1 150
bgneal@1 151 g.remove_edge(e)
bgneal@1 152 pprint.pprint(g)
bgneal@1 153
bgneal@1 154 print "vertices: ", g.vertices()
bgneal@1 155 print "edges: ", g.edges()
bgneal@1 156
bgneal@1 157 g.add_edge(e)
bgneal@1 158 u = Vertex('u')
bgneal@1 159 e1 = Edge(u, v)
bgneal@1 160 e2 = Edge(u, w)
bgneal@1 161 g.add_vertex(u)
bgneal@1 162 g.add_edge(e1)
bgneal@1 163 g.add_edge(e2)
bgneal@1 164 print "Adding vertex u and edges:"
bgneal@1 165 pprint.pprint(g)
bgneal@1 166 print "vertices: ", g.vertices()
bgneal@1 167 print "edges: ", g.edges()
bgneal@1 168
bgneal@1 169 print "Out vertices for v: ", g.out_vertices(v)
bgneal@1 170 print "Out edges for v: ", g.out_edges(v)
bgneal@1 171
bgneal@1 172 x = Vertex('x')
bgneal@1 173 g.add_vertex(x)
bgneal@1 174 g.add_all_edges()
bgneal@1 175 pprint.pprint(g)
bgneal@1 176
bgneal@1 177
bgneal@1 178 if __name__ == '__main__':
bgneal@1 179 import sys
bgneal@1 180 main(*sys.argv)