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