edgedfs.py 5.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174
  1. """
  2. ===========================
  3. Depth First Search on Edges
  4. ===========================
  5. Algorithms for a depth-first traversal of edges in a graph.
  6. """
  7. import networkx as nx
  8. FORWARD = "forward"
  9. REVERSE = "reverse"
  10. __all__ = ["edge_dfs"]
  11. def edge_dfs(G, source=None, orientation=None):
  12. """A directed, depth-first-search of edges in `G`, beginning at `source`.
  13. Yield the edges of G in a depth-first-search order continuing until
  14. all edges are generated.
  15. Parameters
  16. ----------
  17. G : graph
  18. A directed/undirected graph/multigraph.
  19. source : node, list of nodes
  20. The node from which the traversal begins. If None, then a source
  21. is chosen arbitrarily and repeatedly until all edges from each node in
  22. the graph are searched.
  23. orientation : None | 'original' | 'reverse' | 'ignore' (default: None)
  24. For directed graphs and directed multigraphs, edge traversals need not
  25. respect the original orientation of the edges.
  26. When set to 'reverse' every edge is traversed in the reverse direction.
  27. When set to 'ignore', every edge is treated as undirected.
  28. When set to 'original', every edge is treated as directed.
  29. In all three cases, the yielded edge tuples add a last entry to
  30. indicate the direction in which that edge was traversed.
  31. If orientation is None, the yielded edge has no direction indicated.
  32. The direction is respected, but not reported.
  33. Yields
  34. ------
  35. edge : directed edge
  36. A directed edge indicating the path taken by the depth-first traversal.
  37. For graphs, `edge` is of the form `(u, v)` where `u` and `v`
  38. are the tail and head of the edge as determined by the traversal.
  39. For multigraphs, `edge` is of the form `(u, v, key)`, where `key` is
  40. the key of the edge. When the graph is directed, then `u` and `v`
  41. are always in the order of the actual directed edge.
  42. If orientation is not None then the edge tuple is extended to include
  43. the direction of traversal ('forward' or 'reverse') on that edge.
  44. Examples
  45. --------
  46. >>> nodes = [0, 1, 2, 3]
  47. >>> edges = [(0, 1), (1, 0), (1, 0), (2, 1), (3, 1)]
  48. >>> list(nx.edge_dfs(nx.Graph(edges), nodes))
  49. [(0, 1), (1, 2), (1, 3)]
  50. >>> list(nx.edge_dfs(nx.DiGraph(edges), nodes))
  51. [(0, 1), (1, 0), (2, 1), (3, 1)]
  52. >>> list(nx.edge_dfs(nx.MultiGraph(edges), nodes))
  53. [(0, 1, 0), (1, 0, 1), (0, 1, 2), (1, 2, 0), (1, 3, 0)]
  54. >>> list(nx.edge_dfs(nx.MultiDiGraph(edges), nodes))
  55. [(0, 1, 0), (1, 0, 0), (1, 0, 1), (2, 1, 0), (3, 1, 0)]
  56. >>> list(nx.edge_dfs(nx.DiGraph(edges), nodes, orientation="ignore"))
  57. [(0, 1, 'forward'), (1, 0, 'forward'), (2, 1, 'reverse'), (3, 1, 'reverse')]
  58. >>> list(nx.edge_dfs(nx.MultiDiGraph(edges), nodes, orientation="ignore"))
  59. [(0, 1, 0, 'forward'), (1, 0, 0, 'forward'), (1, 0, 1, 'reverse'), (2, 1, 0, 'reverse'), (3, 1, 0, 'reverse')]
  60. Notes
  61. -----
  62. The goal of this function is to visit edges. It differs from the more
  63. familiar depth-first traversal of nodes, as provided by
  64. :func:`~networkx.algorithms.traversal.depth_first_search.dfs_edges`, in
  65. that it does not stop once every node has been visited. In a directed graph
  66. with edges [(0, 1), (1, 2), (2, 1)], the edge (2, 1) would not be visited
  67. if not for the functionality provided by this function.
  68. See Also
  69. --------
  70. :func:`~networkx.algorithms.traversal.depth_first_search.dfs_edges`
  71. """
  72. nodes = list(G.nbunch_iter(source))
  73. if not nodes:
  74. return
  75. directed = G.is_directed()
  76. kwds = {"data": False}
  77. if G.is_multigraph() is True:
  78. kwds["keys"] = True
  79. # set up edge lookup
  80. if orientation is None:
  81. def edges_from(node):
  82. return iter(G.edges(node, **kwds))
  83. elif not directed or orientation == "original":
  84. def edges_from(node):
  85. for e in G.edges(node, **kwds):
  86. yield e + (FORWARD,)
  87. elif orientation == "reverse":
  88. def edges_from(node):
  89. for e in G.in_edges(node, **kwds):
  90. yield e + (REVERSE,)
  91. elif orientation == "ignore":
  92. def edges_from(node):
  93. for e in G.edges(node, **kwds):
  94. yield e + (FORWARD,)
  95. for e in G.in_edges(node, **kwds):
  96. yield e + (REVERSE,)
  97. else:
  98. raise nx.NetworkXError("invalid orientation argument.")
  99. # set up formation of edge_id to easily look up if edge already returned
  100. if directed:
  101. def edge_id(edge):
  102. # remove direction indicator
  103. return edge[:-1] if orientation is not None else edge
  104. else:
  105. def edge_id(edge):
  106. # single id for undirected requires frozenset on nodes
  107. return (frozenset(edge[:2]),) + edge[2:]
  108. # Basic setup
  109. check_reverse = directed and orientation in ("reverse", "ignore")
  110. visited_edges = set()
  111. visited_nodes = set()
  112. edges = {}
  113. # start DFS
  114. for start_node in nodes:
  115. stack = [start_node]
  116. while stack:
  117. current_node = stack[-1]
  118. if current_node not in visited_nodes:
  119. edges[current_node] = edges_from(current_node)
  120. visited_nodes.add(current_node)
  121. try:
  122. edge = next(edges[current_node])
  123. except StopIteration:
  124. # No more edges from the current node.
  125. stack.pop()
  126. else:
  127. edgeid = edge_id(edge)
  128. if edgeid not in visited_edges:
  129. visited_edges.add(edgeid)
  130. # Mark the traversed "to" node as to-be-explored.
  131. if check_reverse and edge[-1] == REVERSE:
  132. stack.append(edge[0])
  133. else:
  134. stack.append(edge[1])
  135. yield edge