1""" 2================== 3Depth-first search 4================== 5 6Basic algorithms for depth-first searching. 7 8Based on http://www.ics.uci.edu/~eppstein/PADS/DFS.py 9by D. Eppstein, July 2004. 10""" 11__author__ = """\n""".join(['Aric Hagberg <hagberg@lanl.gov>']) 12 13__all__ = ['dfs_edges', 'dfs_tree', 14 'dfs_predecessors', 'dfs_successors', 15 'dfs_preorder_nodes','dfs_postorder_nodes', 16 'dfs_labeled_edges'] 17 18import networkx as nx 19from collections import defaultdict 20 21def dfs_edges(G,source=None): 22 """Produce edges in a depth-first-search starting at source.""" 23 # Based on http://www.ics.uci.edu/~eppstein/PADS/DFS.py 24 # by D. Eppstein, July 2004. 25 if source is None: 26 # produce edges for all components 27 nodes=G 28 else: 29 # produce edges for components with source 30 nodes=[source] 31 visited=set() 32 for start in nodes: 33 if start in visited: 34 continue 35 visited.add(start) 36 stack = [(start,iter(G[start]))] 37 while stack: 38 parent,children = stack[-1] 39 try: 40 child = next(children) 41 if child not in visited: 42 yield parent,child 43 visited.add(child) 44 stack.append((child,iter(G[child]))) 45 except StopIteration: 46 stack.pop() 47 48def dfs_tree(G, source): 49 """Return directed tree of depth-first-search from source.""" 50 T = nx.DiGraph() 51 if source is None: 52 T.add_nodes_from(G) 53 else: 54 T.add_node(source) 55 T.add_edges_from(dfs_edges(G,source)) 56 return T 57 58def dfs_predecessors(G, source=None): 59 """Return dictionary of predecessors in depth-first-search from source.""" 60 return dict((t,s) for s,t in dfs_edges(G,source=source)) 61 62 63def dfs_successors(G, source=None): 64 """Return dictionary of successors in depth-first-search from source.""" 65 d=defaultdict(list) 66 for s,t in dfs_edges(G,source=source): 67 d[s].append(t) 68 return dict(d) 69 70 71def dfs_postorder_nodes(G,source=None): 72 """Produce nodes in a depth-first-search post-ordering starting 73 from source. 74 """ 75 post=(v for u,v,d in nx.dfs_labeled_edges(G,source=source) 76 if d['dir']=='reverse') 77 # chain source to end of pre-ordering 78# return chain(post,[source]) 79 return post 80 81 82def dfs_preorder_nodes(G,source=None): 83 """Produce nodes in a depth-first-search pre-ordering starting at source.""" 84 pre=(v for u,v,d in nx.dfs_labeled_edges(G,source=source) 85 if d['dir']=='forward') 86 # chain source to beginning of pre-ordering 87# return chain([source],pre) 88 return pre 89 90 91def dfs_labeled_edges(G,source=None): 92 """Produce edges in a depth-first-search starting at source and 93 labeled by direction type (forward, reverse, nontree). 94 """ 95 # Based on http://www.ics.uci.edu/~eppstein/PADS/DFS.py 96 # by D. Eppstein, July 2004. 97 if source is None: 98 # produce edges for all components 99 nodes=G 100 else: 101 # produce edges for components with source 102 nodes=[source] 103 visited=set() 104 for start in nodes: 105 if start in visited: 106 continue 107 yield start,start,{'dir':'forward'} 108 visited.add(start) 109 stack = [(start,iter(G[start]))] 110 while stack: 111 parent,children = stack[-1] 112 try: 113 child = next(children) 114 if child in visited: 115 yield parent,child,{'dir':'nontree'} 116 else: 117 yield parent,child,{'dir':'forward'} 118 visited.add(child) 119 stack.append((child,iter(G[child]))) 120 except StopIteration: 121 stack.pop() 122 if stack: 123 yield stack[-1][0],parent,{'dir':'reverse'} 124 yield start,start,{'dir':'reverse'} 125