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