1  
  2  class Graph:
  3      def __init__(self, n, edges, directed=False):
  4          """
  5          :param n: The number of vertices
  6          :param edges: an iterable of (u, v) edge tuples
  7          :param directed: boolean, whether this graph is directed
  8          """
  9          self.edges = [set() for _ in range(n)]
 10          for u, v in edges:
 11              self.edges[u].add(v)
 12              if not directed:
 13                  self.edges[v].add(u)
 14  
 15  
 16      def is_edge(self, u, v):
 17          """
 18          :return: Returns whether node (u, v) is present
 19          """
 20          return v in self.edges[u]
 21  
 22  
 23      def neighbors(self, u):
 24          """
 25          :return: A set of u's neighbors
 26          """
 27          return self.edges[u].copy()
 28  
 29  
 30      def size(self):
 31          """
 32          :return: The number of vertices
 33          """
 34          return len(self.edges)
 35  
 36  
 37  from collections import deque
 38  
 39  """
 40  BFS on graph using a queue. Returns dictionaries of parents and
 41  distance from source for each node in BFS tree
 42  """
 43  def bfs(graph: Graph, source: int):
 44      parent = {source: source}
 45      dist = {source: 0}
 46      queue = deque([source])
 47  
 48      while len(queue) != 0:
 49          u = queue.popleft()
 50          for v in graph.neighbors(u):
 51              if v not in dist:
 52                  parent[v] = u
 53                  dist[v] = dist[u] + 1
 54                  queue.append(v)
 55  
 56      return parent, dist
 57  
 58  
 59  """
 60  Recursive DFS: returns dictionaries of start time, end time,
 61  and parents for each node in DFS forest
 62  """
 63  def dfs_re(graph: Graph, source: int):
 64      start, end, parent, time = {}, {}, {}, [1]
 65  
 66      def dfs_visit(u):
 67          start[u] = time[0]
 68          time[0] += 1
 69          for v in graph.neighbors(u):
 70              if v not in start:
 71                  parent[v] = u
 72                  dfs_visit(v)
 73          end[u] = time[0]
 74          time[0] += 1
 75  
 76      dfs_visit(source)
 77  
 78      for u in range(graph.size()):
 79          if u not in start:
 80              parent[u] = u
 81              dfs_visit(u)
 82  
 83      return parent, start, end
 84  
 85  
 86  """
 87  Iterative DFS: uses stack, returns dictionary of node -> parent
 88  """
 89  def dfs_it(graph: Graph, source: int):
 90      stack = deque([source])
 91      parent = {source: source}
 92  
 93      while len(stack) != 0:
 94          u = stack.pop()
 95          for v in graph.neighbors(u):
 96              if v not in parent:
 97                  parent[v] = u
 98                  stack.append(v)
 99  
100      return parent
101  
102  
103  edges = [()]
104  g = Graph(6, [(0, 1), (0, 2), (1, 2), (1, 3), (2, 4), (4, 5)])
105  
106  # p, d = bfs(g, 0)
107  # print(d)
108  
109  p = dfs_it(g, 0)
110  print(p)