3 # © Copyright 2021-2022, Scott Gasch
5 """A simple graph class that can be optionally directed and weighted and
6 some operations on it."""
9 from typing import Dict, Generator, List, Optional, Set
11 from pyutils.types.simple import Numeric
15 def __init__(self, directed: bool = False):
16 """Constructs a new Graph object.
19 directed: are we modeling a directed graph? See :meth:`add_edge`.
22 self.directed = directed
23 self.graph: Dict[str, Dict[str, Numeric]] = {}
25 def add_vertex(self, vertex_id: str) -> bool:
26 """Adds a new vertex to the graph.
29 vertex_id: the unique identifier of the new vertex.
32 True unless vertex_id is already in the graph.
41 >>> len(g.get_vertices())
45 if vertex_id not in self.graph:
46 self.graph[vertex_id] = {}
50 def add_edge(self, src: str, dest: str, weight: Numeric = 1) -> None:
51 """Adds a new (optionally weighted) edge between src and dest
52 vertexes. If the graph is not directed (see c'tor) this also
53 adds a reciprocal edge with the same weight back from dest to
58 If either or both of src and dest are not already added to
59 the graph, they are implicitly added by adding this edge.
62 src: the source vertex id
63 dest: the destination vertex id
64 weight: optionally, the weight of the edge(s) added
67 >>> g.add_edge('a', 'b')
68 >>> g.add_edge('b', 'c', weight=2)
69 >>> len(g.get_vertices())
72 {'a': {'b': 1}, 'b': {'a': 1, 'c': 2}, 'c': {'b': 2}}
77 self.graph[src][dest] = weight
79 self.graph[dest][src] = weight
81 def get_vertices(self) -> List[str]:
84 a list of the vertex ids in the graph.
89 >>> g.add_edge('b', 'c')
93 return list(self.graph.keys())
95 def get_edges(self) -> Dict[str, Dict[str, Numeric]]:
98 A dict whose keys are source vertexes and values
99 are dicts of destination vertexes with values describing the
100 weight of the edge from source to destination.
102 >>> g = Graph(directed=True)
103 >>> g.add_edge('a', 'b')
104 >>> g.add_edge('b', 'c', weight=2)
105 >>> len(g.get_vertices())
108 {'a': {'b': 1}, 'b': {'c': 2}, 'c': {}}
112 def _dfs(self, vertex: str, visited: Set[str]):
115 for neighbor in self.graph[vertex]:
116 if neighbor not in visited:
117 yield from self._dfs(neighbor, visited)
120 self, starting_vertex: str, target: Optional[str] = None
121 ) -> Generator[str, None, None]:
122 """Performs a depth first traversal of the graph.
125 starting_vertex: The DFS starting point.
126 target: The vertex that, if found, indicates to halt.
129 An ordered sequence of vertex ids visited by the traversal.
136 A -- C -- D -- E -- F;
142 >>> g.add_edge('A', 'B')
143 >>> g.add_edge('A', 'C')
144 >>> g.add_edge('B', 'D')
145 >>> g.add_edge('C', 'D')
146 >>> g.add_edge('D', 'E')
147 >>> g.add_edge('E', 'F')
148 >>> g.add_edge('E', 'G')
149 >>> g.add_edge('F', 'F')
150 >>> for node in g.dfs('A'):
160 >>> for node in g.dfs('F', 'B'):
167 visited: Set[str] = set()
168 for node in self._dfs(starting_vertex, visited):
174 self, starting_vertex: str, target: Optional[str] = None
175 ) -> Generator[str, None, None]:
176 """Performs a breadth first traversal of the graph.
179 starting_vertex: The BFS starting point.
180 target: The vertex that, if found, we should halt the search.
183 An ordered sequence of vertex ids visited by the traversal.
190 A -- C -- D -- E -- F;
196 >>> g.add_edge('A', 'B')
197 >>> g.add_edge('A', 'C')
198 >>> g.add_edge('B', 'D')
199 >>> g.add_edge('C', 'D')
200 >>> g.add_edge('D', 'E')
201 >>> g.add_edge('E', 'F')
202 >>> g.add_edge('E', 'G')
203 >>> g.add_edge('F', 'F')
204 >>> for node in g.bfs('A'):
214 >>> for node in g.bfs('F', 'G'):
224 todo.append(starting_vertex)
225 visited.add(starting_vertex)
233 neighbors = self.graph[vertex]
234 for neighbor in neighbors:
235 if neighbor not in visited:
236 todo.append(neighbor)
237 visited.add(neighbor)
240 if __name__ == "__main__":