Mercurial > public > think_complexity

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