fix digraph_tree

This commit is contained in:
Costa Shulyupin 2018-07-29 16:03:45 +03:00
parent f6f4513543
commit 75da8dd37a
1 changed files with 5 additions and 4 deletions

View File

@ -416,14 +416,16 @@ def starts(dg): # roots
return {n: dg.out_degree(n) for (n, d) in dg.in_degree if not d} return {n: dg.out_degree(n) for (n, d) in dg.in_degree if not d}
def digraph_tree(dg, starts=None): def digraph_tree(dg, starts=None, black_list=black_list):
tree = nx.DiGraph() tree = nx.DiGraph()
def sub(node): def sub(node):
tree.add_node(node) tree.add_node(node)
for o in dg.successors(node): for o in dg.successors(node):
if o in black_list or tree.has_node(o): if o in black_list or tree.has_edge(node, o) or o in starts:
# print(o)
continue continue
tree.add_edge(node,o) tree.add_edge(node, o)
sub(o) sub(o)
printed = set() printed = set()
@ -438,7 +440,6 @@ def digraph_tree(dg, starts=None):
for o in starts: for o in starts:
if o in black_list: if o in black_list:
continue continue
tree.add_edge(default_root,o)
sub(o) sub(o)
return tree return tree