src/pyutils/graph.py

   1 #!/usr/bin/env python3
   2
   3 # © Copyright 2021-2022, Scott Gasch
   4
   5 """A simple graph class that can be optionally directed and weighted and
   6 some operations on it."""
   7
   8
   9 from typing import Dict, Generator, List, Optional, Set
  10
  11 from pyutils.types.simple import Numeric
  12
  13
  14 class Graph(object):
  15     def __init__(self, directed: bool = False):
  16         """Constructs a new Graph object.
  17
  18         Args:
  19             directed: are we modeling a directed graph?  See ::meth
  20                 add_edge.
  21
  22         """
  23         self.directed = directed
  24         self.graph: Dict[str, Dict[str, Numeric]] = {}
  25
  26     def add_vertex(self, vertex_id: str) -> bool:
  27         """Adds a new vertex to the graph.
  28
  29         Args:
  30             vertex_id: the unique identifier of the new vertex.
  31
  32         Returns:
  33             True unless vertex_id is already in the graph.
  34
  35         >>> g = Graph()
  36         >>> g.add_vertex('a')
  37         True
  38         >>> g.add_vertex('b')
  39         True
  40         >>> g.add_vertex('a')
  41         False
  42         >>> len(g.get_vertices())
  43         2
  44
  45         """
  46         if vertex_id not in self.graph:
  47             self.graph[vertex_id] = {}
  48             return True
  49         return False
  50
  51     def add_edge(self, src: str, dest: str, weight: Numeric = 1) -> None:
  52         """Adds a new (optionally weighted) edge between src and dest
  53         vertexes.  If the graph is not directed (see c'tor) this also
  54         adds a reciprocal edge with the same weight back from dest to
  55         src too.
  56
  57         .. note::
  58
  59         If either or both of src and dest are not already added to the
  60         graph, they are implicitly added by adding this edge.
  61
  62         Args:
  63             src: the source vertex id
  64             dest: the destination vertex id
  65             weight: optionally, the weight of the edge(s) added
  66
  67         >>> g = Graph()
  68         >>> g.add_edge('a', 'b')
  69         >>> g.add_edge('b', 'c', weight=2)
  70         >>> len(g.get_vertices())
  71         3
  72         >>> g.get_edges()
  73         {'a': {'b': 1}, 'b': {'a': 1, 'c': 2}, 'c': {'b': 2}}
  74
  75         """
  76         self.add_vertex(src)
  77         self.add_vertex(dest)
  78         self.graph[src][dest] = weight
  79         if not self.directed:
  80             self.graph[dest][src] = weight
  81
  82     def get_vertices(self) -> List[str]:
  83         """
  84         Returns:
  85             a list of the vertex ids in the graph.
  86
  87         >>> g = Graph()
  88         >>> g.add_vertex('a')
  89         True
  90         >>> g.add_edge('b', 'c')
  91         >>> g.get_vertices()
  92         ['a', 'b', 'c']
  93         """
  94         return list(self.graph.keys())
  95
  96     def get_edges(self) -> Dict[str, Dict[str, Numeric]]:
  97         """
  98         Returns:
  99             A dict whose keys are source vertexes and values
 100             are dicts of destination vertexes with values describing the
 101             weight of the edge from source to destination.
 102
 103         >>> g = Graph(directed=True)
 104         >>> g.add_edge('a', 'b')
 105         >>> g.add_edge('b', 'c', weight=2)
 106         >>> len(g.get_vertices())
 107         3
 108         >>> g.get_edges()
 109         {'a': {'b': 1}, 'b': {'c': 2}, 'c': {}}
 110         """
 111         return self.graph
 112
 113     def _dfs(self, vertex: str, visited: Set[str]):
 114         yield vertex
 115         visited.add(vertex)
 116         for neighbor in self.graph[vertex]:
 117             if neighbor not in visited:
 118                 yield from self._dfs(neighbor, visited)
 119
 120     def dfs(
 121         self, starting_vertex: str, target: Optional[str] = None
 122     ) -> Generator[str, None, None]:
 123         """Performs a depth first traversal of the graph.
 124
 125         Args:
 126             starting_vertex: The DFS starting point.
 127             target: The vertex that, if found, indicates to halt.
 128
 129         Returns:
 130             An ordered sequence of vertex ids visited by the traversal.
 131
 132                 A ------ B
 133                 |        |
 134                 |        |
 135                 C ------ D ------ E ------ F -O
 136                                   |
 137                                   |
 138                                   G
 139
 140         >>> g = Graph()
 141         >>> g.add_edge('A', 'B')
 142         >>> g.add_edge('A', 'C')
 143         >>> g.add_edge('B', 'D')
 144         >>> g.add_edge('C', 'D')
 145         >>> g.add_edge('D', 'E')
 146         >>> g.add_edge('E', 'F')
 147         >>> g.add_edge('E', 'G')
 148         >>> g.add_edge('F', 'F')
 149         >>> for node in g.dfs('A'):
 150         ...     print(node)
 151         A
 152         B
 153         D
 154         C
 155         E
 156         F
 157         G
 158
 159         >>> for node in g.dfs('F', 'B'):
 160         ...     print(node)
 161         F
 162         E
 163         D
 164         B
 165         """
 166         visited: Set[str] = set()
 167         for node in self._dfs(starting_vertex, visited):
 168             yield node
 169             if node == target:
 170                 return
 171
 172     def bfs(
 173         self, starting_vertex: str, target: Optional[str] = None
 174     ) -> Generator[str, None, None]:
 175         """Performs a breadth first traversal of the graph.
 176
 177         Args:
 178             starting_vertex: The BFS starting point.
 179             target: The vertex that, if found, we should halt the search.
 180
 181         Returns:
 182             An ordered sequence of vertex ids visited by the traversal.
 183
 184                 A ------ B
 185                 |        |
 186                 |        |
 187                 C ------ D ------ E ------ F -O
 188                                   |
 189                                   |
 190                                   G
 191
 192         >>> g = Graph()
 193         >>> g.add_edge('A', 'B')
 194         >>> g.add_edge('A', 'C')
 195         >>> g.add_edge('B', 'D')
 196         >>> g.add_edge('C', 'D')
 197         >>> g.add_edge('D', 'E')
 198         >>> g.add_edge('E', 'F')
 199         >>> g.add_edge('E', 'G')
 200         >>> g.add_edge('F', 'F')
 201         >>> for node in g.bfs('A'):
 202         ...     print(node)
 203         A
 204         B
 205         C
 206         D
 207         E
 208         F
 209         G
 210
 211         >>> for node in g.bfs('F', 'G'):
 212         ...     print(node)
 213         F
 214         E
 215         D
 216         G
 217         """
 218         todo = []
 219         visited = set()
 220
 221         todo.append(starting_vertex)
 222         visited.add(starting_vertex)
 223
 224         while todo:
 225             vertex = todo.pop(0)
 226             yield vertex
 227             if vertex == target:
 228                 return
 229
 230             neighbors = self.graph[vertex]
 231             for neighbor in neighbors:
 232                 if neighbor not in visited:
 233                     todo.append(neighbor)
 234                     visited.add(neighbor)
 235
 236
 237 if __name__ == "__main__":
 238     import doctest
 239
 240     doctest.testmod()