1from __future__ import generators 2 3class DirectedGraph(object): 4 '''This class is for directed graphs with vertices of arbitrary type''' 5 def __init__(self, vertices = []): 6 '''Create a graph''' 7 self.vertices = [] 8 self.inEdges = {} 9 self.outEdges = {} 10 map(self.addVertex, vertices) 11 return 12 13 def __len__(self): 14 return len(self.vertices) 15 16 def __str__(self): 17 return 'DirectedGraph with '+str(len(self.vertices))+' vertices and '+str(reduce(lambda k,l: k+l, [len(edgeList) for edgeList in self.inEdges.values()], 0))+' edges' 18 19 def addVertex(self, vertex): 20 '''Add a vertex if it does not already exist in the vertex list 21 - Should be able to use Set in Python 2.3''' 22 if vertex is None: return 23 if not vertex in self.vertices: 24 self.vertices.append(vertex) 25 self.clearEdges(vertex) 26 return 27 28 def addEdges(self, vertex, inputs = [], outputs = []): 29 '''Define the in and out edges for a vertex by listing the other vertices defining the edges 30 - If any vertex does not exist in the graph, it is created''' 31 self.addVertex(vertex) 32 for input in inputs: 33 self.addVertex(input) 34 if not vertex is None and not input is None: 35 if not input in self.inEdges[vertex]: self.inEdges[vertex].append(input) 36 if not vertex in self.outEdges[input]: self.outEdges[input].append(vertex) 37 for output in outputs: 38 self.addVertex(output) 39 if not vertex is None and not output is None: 40 if not vertex in self.inEdges[output]: self.inEdges[output].append(vertex) 41 if not output in self.outEdges[vertex]: self.outEdges[vertex].append(output) 42 return 43 44 def getEdges(self, vertex): 45 return (self.inEdges[vertex], self.outEdges[vertex]) 46 47 def clearEdges(self, vertex, inOnly = 0, outOnly = 0): 48 if inOnly and outOnly: 49 raise RuntimeError('Inconsistent arguments') 50 if not outOnly: 51 self.inEdges[vertex] = [] 52 if not inOnly: 53 self.outEdges[vertex] = [] 54 return 55 56 def removeVertex(self, vertex): 57 '''Remove a vertex if it already exists in the vertex list 58 - Also removes all associated edges''' 59 if vertex is None: return 60 if vertex in self.vertices: 61 self.vertices.remove(vertex) 62 del self.inEdges[vertex] 63 del self.outEdges[vertex] 64 for v in self.vertices: 65 if vertex in self.inEdges[v]: self.inEdges[v].remove(vertex) 66 if vertex in self.outEdges[v]: self.outEdges[v].remove(vertex) 67 return 68 69 def replaceVertex(self, vertex, newVertex): 70 '''Replace a vertex with newVertex if it already exists in the vertex list 71 - Also transfers all associated edges''' 72 if vertex is None or newVertex is None: return 73 self.addEdges(newVertex, self.inEdges[vertex], self.outEdges[vertex]) 74 self.removeVertex(vertex) 75 return 76 77 def addSubgraph(self, graph): 78 '''Add the vertices and edges of another graph into this one''' 79 map(self.addVertex, graph.vertices) 80 map(lambda v: apply(self.addEdges, (v,)+graph.getEdges(v)), graph.vertices) 81 return 82 83 def removeSubgraph(self, graph): 84 '''Remove the vertices and edges of a subgraph, and all the edges connected to it''' 85 map(self.removeVertex, graph.vertices) 86 return 87 88 def printIndent(self, indent): 89 import sys 90 for i in range(indent): sys.stdout.write(' ') 91 92 def display(self): 93 print 'I am a DirectedGraph with '+str(len(self.vertices))+' vertices' 94 for vertex in DirectedGraph.breadthFirstSearch(self): 95 self.printIndent(vertex.__level) 96 print '('+str(self.vertices.index(vertex))+') '+str(vertex)+' in: '+str(map(self.vertices.index, self.inEdges[vertex]))+' out: '+str(map(self.vertices.index, self.outEdges[vertex])) 97 return 98 99 def appendGraph(self, graph): 100 '''Join every leaf of this graph to every root of the input graph, leaving the result in this graph''' 101 leaves = DirectedGraph.getLeaves(self) 102 self.addSubgraph(graph) 103 map(lambda v: self.addEdges(v, outputs = DirectedGraph.getRoots(graph)), leaves) 104 return self 105 106 def prependGraph(self, graph): 107 '''Join every leaf of the input graph to every root of this graph, leaving the result in this graph''' 108 roots = DirectedGraph.getRoots(self) 109 self.addSubgraph(graph) 110 map(lambda v: self.addEdges(v, outputs = roots), DirectedGraph.getLeaves(graph)) 111 return self 112 113 def getRoots(graph): 114 '''Return all the sources in the graph (nodes without entering edges)''' 115 return filter(lambda v: not len(graph.getEdges(v)[0]), graph.vertices) 116 getRoots = staticmethod(getRoots) 117 118 def getLeaves(graph): 119 '''Return all the sinks in the graph (nodes without exiting edges)''' 120 return filter(lambda v: not len(graph.getEdges(v)[1]), graph.vertices) 121 getLeaves = staticmethod(getLeaves) 122 123 def depthFirstVisit(graph, vertex, seen = None, returnFinished = 0, outEdges = 1): 124 '''This is a generator returning vertices in a depth-first traversal only for the subtree rooted at vertex 125 - If returnFinished is True, return a vertex when it finishes 126 - Otherwise, return a vertex when it is first seen 127 - If outEdges is True, proceed along these, otherwise use inEdges''' 128 if seen is None: seen = [] 129 seen.append(vertex) 130 if not returnFinished: 131 yield vertex 132 # Cute trick since outEdges is index 1, and inEdges is index 0 133 for v in graph.getEdges(vertex)[outEdges]: 134 if not v in seen: 135 try: 136 for v2 in DirectedGraph.depthFirstVisit(graph, v, seen, returnFinished, outEdges): 137 yield v2 138 except StopIteration: 139 pass 140 if returnFinished: 141 yield vertex 142 return 143 depthFirstVisit = staticmethod(depthFirstVisit) 144 145 def depthFirstSearch(graph, returnFinished = 0, outEdges = 1): 146 '''This is a generator returning vertices in a depth-first traversal 147 - If returnFinished is True, return a vertex when it finishes 148 - Otherwise, return a vertex when it is first seen 149 - If outEdges is True, proceed along these, otherwise use inEdges''' 150 seen = [] 151 for vertex in graph.vertices: 152 if not vertex in seen: 153 try: 154 for v in DirectedGraph.depthFirstVisit(graph, vertex, seen, returnFinished, outEdges): 155 yield v 156 except StopIteration: 157 pass 158 return 159 depthFirstSearch = staticmethod(depthFirstSearch) 160 161 def breadthFirstSearch(graph, returnFinished = 0): 162 '''This is a generator returning vertices in a breadth-first traversal 163 - If returnFinished is True, return a vertex when it finishes 164 - Otherwise, return a vertex when it is first seen''' 165 queue = DirectedGraph.getRoots(graph)[0:1] 166 if not len(queue): return 167 seen = [queue[0]] 168 if not returnFinished: 169 queue[0].__level = 0 170 yield queue[0] 171 while len(queue): 172 vertex = queue[-1] 173 for v in graph.getEdges(vertex)[1]: 174 if not v in seen: 175 seen.append(v) 176 v.__level = vertex.__level + 1 177 queue.insert(0, v) 178 if not returnFinished: 179 yield v 180 vertex = queue.pop() 181 if returnFinished: 182 yield vertex 183 return 184 185 def topologicalSort(graph, start = None, outEdges = 1): 186 '''Reorder the vertices using topological sort''' 187 if start is None: 188 vertices = [vertex for vertex in DirectedGraph.depthFirstSearch(graph, returnFinished = 1, outEdges = outEdges)] 189 else: 190 vertices = [vertex for vertex in DirectedGraph.depthFirstVisit(graph, start, returnFinished = 1, outEdges = outEdges)] 191 vertices.reverse() 192 for vertex in vertices: 193 yield vertex 194 return 195 topologicalSort = staticmethod(topologicalSort) 196