-
Notifications
You must be signed in to change notification settings - Fork 240
/
graph_utils.py
53 lines (45 loc) · 1.2 KB
/
graph_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
"""
Various graph related utilities.
"""
import networkx as nx
def get_sccs_topo(graph):
"""
Return strongly connected subsystems of the given Group in topological order.
Parameters
----------
graph : networkx.DiGraph
Directed graph of Systems.
Returns
-------
list of sets of str
A list of strongly connected components in topological order.
"""
# Tarjan's algorithm returns SCCs in reverse topological order, so
# the list returned here is reversed.
sccs = list(nx.strongly_connected_components(graph))
sccs.reverse()
return sccs
def all_connected_nodes(graph, start):
"""
Yield all downstream nodes starting at the given node.
Parameters
----------
graph : network.DiGraph
Graph being traversed.
start : hashable object
Identifier of the starting node.
Yields
------
str
Each node found when traversal starts at start.
"""
stack = [start]
visited = set(stack)
yield start
while stack:
src = stack.pop()
for tgt in graph[src]:
yield tgt
if tgt not in visited:
visited.add(tgt)
stack.append(tgt)