1"""
2====================
3Breadth-first search
4====================
5
6Basic algorithms for breadth-first searching.
7"""
8__author__ = """\n""".join(['Aric Hagberg <hagberg@lanl.gov>'])
9
10__all__ = ['bfs_edges', 'bfs_tree',
11           'bfs_predecessors', 'bfs_successors']
12
13import networkx as nx
14from collections import defaultdict, deque
15
16def bfs_edges(G, source, reverse=False):
17    """Produce edges in a breadth-first-search starting at source."""
18    # Based on http://www.ics.uci.edu/~eppstein/PADS/BFS.py
19    # by D. Eppstein, July 2004.
20    if reverse and isinstance(G, nx.DiGraph):
21        neighbors = G.predecessors_iter
22    else:
23        neighbors = G.neighbors_iter
24    visited=set([source])
25    queue = deque([(source, neighbors(source))])
26    while queue:
27        parent, children = queue[0]
28        try:
29            child = next(children)
30            if child not in visited:
31                yield parent, child
32                visited.add(child)
33                queue.append((child, neighbors(child)))
34        except StopIteration:
35            queue.popleft()
36
37def bfs_tree(G, source, reverse=False):
38    """Return directed tree of breadth-first-search from source."""
39    T = nx.DiGraph()
40    T.add_node(source)
41    T.add_edges_from(bfs_edges(G,source,reverse=reverse))
42    return T
43
44def bfs_predecessors(G, source):
45    """Return dictionary of predecessors in breadth-first-search from source."""
46    return dict((t,s) for s,t in bfs_edges(G,source))
47
48def bfs_successors(G, source):
49    """Return dictionary of successors in breadth-first-search from source."""
50    d=defaultdict(list)
51    for s,t in bfs_edges(G,source):
52        d[s].append(t)
53    return dict(d)
54