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:`add_edge`.
  20
  21         """
  22         self.directed = directed
  23         self.graph: Dict[str, Dict[str, Numeric]] = {}
  24
  25     def add_vertex(self, vertex_id: str) -> bool:
  26         """Adds a new vertex to the graph.
  27
  28         Args:
  29             vertex_id: the unique identifier of the new vertex.
  30
  31         Returns:
  32             True unless vertex_id is already in the graph.
  33
  34         >>> g = Graph()
  35         >>> g.add_vertex('a')
  36         True
  37         >>> g.add_vertex('b')
  38         True
  39         >>> g.add_vertex('a')
  40         False
  41         >>> len(g.get_vertices())
  42         2
  43
  44         """
  45         if vertex_id not in self.graph:
  46             self.graph[vertex_id] = {}
  47             return True
  48         return False
  49
  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
  54         src too.
  55
  56         .. note::
  57
  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.
  60
  61         Args:
  62             src: the source vertex id
  63             dest: the destination vertex id
  64             weight: optionally, the weight of the edge(s) added
  65
  66         >>> g = Graph()
  67         >>> g.add_edge('a', 'b')
  68         >>> g.add_edge('b', 'c', weight=2)
  69         >>> len(g.get_vertices())
  70         3
  71         >>> g.get_edges()
  72         {'a': {'b': 1}, 'b': {'a': 1, 'c': 2}, 'c': {'b': 2}}
  73
  74         """
  75         self.add_vertex(src)
  76         self.add_vertex(dest)
  77         self.graph[src][dest] = weight
  78         if not self.directed:
  79             self.graph[dest][src] = weight
  80
  81     def get_vertices(self) -> List[str]:
  82         """
  83         Returns:
  84             a list of the vertex ids in the graph.
  85
  86         >>> g = Graph()
  87         >>> g.add_vertex('a')
  88         True
  89         >>> g.add_edge('b', 'c')
  90         >>> g.get_vertices()
  91         ['a', 'b', 'c']
  92         """
  93         return list(self.graph.keys())
  94
  95     def get_edges(self) -> Dict[str, Dict[str, Numeric]]:
  96         """
  97         Returns:
  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.
 101
 102         >>> g = Graph(directed=True)
 103         >>> g.add_edge('a', 'b')
 104         >>> g.add_edge('b', 'c', weight=2)
 105         >>> len(g.get_vertices())
 106         3
 107         >>> g.get_edges()
 108         {'a': {'b': 1}, 'b': {'c': 2}, 'c': {}}
 109         """
 110         return self.graph
 111
 112     def _dfs(self, vertex: str, visited: Set[str]):
 113         yield vertex
 114         visited.add(vertex)
 115         for neighbor in self.graph[vertex]:
 116             if neighbor not in visited:
 117                 yield from self._dfs(neighbor, visited)
 118
 119     def dfs(
 120         self, starting_vertex: str, target: Optional[str] = None
 121     ) -> Generator[str, None, None]:
 122         """Performs a depth first traversal of the graph.
 123
 124         Args:
 125             starting_vertex: The DFS starting point.
 126             target: The vertex that, if found, indicates to halt.
 127
 128         Returns:
 129             An ordered sequence of vertex ids visited by the traversal.
 130
 131         .. graphviz::
 132
 133             graph g {
 134                 node [shape=record];
 135                 A -- B -- D;
 136                 A -- C -- D -- E -- F;
 137                 F -- F;
 138                 E -- G;
 139             }
 140
 141         >>> g = Graph()
 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'):
 151         ...     print(node)
 152         A
 153         B
 154         D
 155         C
 156         E
 157         F
 158         G
 159
 160         >>> for node in g.dfs('F', 'B'):
 161         ...     print(node)
 162         F
 163         E
 164         D
 165         B
 166         """
 167         visited: Set[str] = set()
 168         for node in self._dfs(starting_vertex, visited):
 169             yield node
 170             if node == target:
 171                 return
 172
 173     def bfs(
 174         self, starting_vertex: str, target: Optional[str] = None
 175     ) -> Generator[str, None, None]:
 176         """Performs a breadth first traversal of the graph.
 177
 178         Args:
 179             starting_vertex: The BFS starting point.
 180             target: The vertex that, if found, we should halt the search.
 181
 182         Returns:
 183             An ordered sequence of vertex ids visited by the traversal.
 184
 185         .. graphviz::
 186
 187             graph g {
 188                 node [shape=record];
 189                 A -- B -- D;
 190                 A -- C -- D -- E -- F;
 191                 F -- F;
 192                 E -- G;
 193             }
 194
 195         >>> g = Graph()
 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'):
 205         ...     print(node)
 206         A
 207         B
 208         C
 209         D
 210         E
 211         F
 212         G
 213
 214         >>> for node in g.bfs('F', 'G'):
 215         ...     print(node)
 216         F
 217         E
 218         D
 219         G
 220         """
 221         todo = []
 222         visited = set()
 223
 224         todo.append(starting_vertex)
 225         visited.add(starting_vertex)
 226
 227         while todo:
 228             vertex = todo.pop(0)
 229             yield vertex
 230             if vertex == target:
 231                 return
 232
 233             neighbors = self.graph[vertex]
 234             for neighbor in neighbors:
 235                 if neighbor not in visited:
 236                     todo.append(neighbor)
 237                     visited.add(neighbor)
 238
 239
 240 if __name__ == "__main__":
 241     import doctest
 242
 243     doctest.testmod()