Skip to content

Commit

Permalink
Fix forest_str to print trees in the right order
Browse files Browse the repository at this point in the history
  • Loading branch information
Erotemic committed Nov 8, 2020
1 parent 40d631d commit 16ac464
Show file tree
Hide file tree
Showing 3 changed files with 103 additions and 101 deletions.
8 changes: 4 additions & 4 deletions networkx/generators/random_graphs.py
Original file line number Diff line number Diff line change
Expand Up @@ -1325,10 +1325,10 @@ def random_ordered_tree(n, seed=None, directed=False):
>>> otree = random_ordered_tree(n=5, seed=3, directed=True)
>>> print(forest_str(otree))
╙── 1
├─╼ 4
│   ├─╼ 3
│   └─╼ 2
└─╼ 0
├─╼ 0
─╼ 4
─╼ 2
└─╼ 3
"""
from networkx.utils import create_py_random_state

Expand Down
152 changes: 75 additions & 77 deletions networkx/readwrite/tests/test_text.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,24 +11,24 @@ def test_directed_tree_str():
node_target = dedent(
"""
╙── 0
├─╼ 2
│   ├─╼ 6
│   └─╼ 5
└─╼ 1
├─╼ 4
└─╼ 3
├─╼ 1
│   ├─╼ 3
│   └─╼ 4
└─╼ 2
├─╼ 5
└─╼ 6
"""
).strip()

label_target = dedent(
"""
╙── node_a
├─╼ node_c
│   ├─╼ node_g
│   └─╼ node_f
└─╼ node_b
├─╼ node_e
└─╼ node_d
├─╼ node_b
│   ├─╼ node_d
│   └─╼ node_e
└─╼ node_c
├─╼ node_f
└─╼ node_g
"""
).strip()

Expand Down Expand Up @@ -63,60 +63,58 @@ def test_empty_graph():
def test_directed_multi_tree_forest():
tree1 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
tree2 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
tree2 = nx.relabel_nodes(tree2, {n: n + len(tree1) for n in tree2.nodes})
forest = nx.union(tree1, tree2)
ret = nx.forest_str(forest, sources=[0, 7])
forest = nx.disjoint_union_all([tree1, tree2])
ret = nx.forest_str(forest)
print(ret)

target = dedent(
"""
╟── 7
╎   ├─╼ 9
╎   │   ├─╼ 13
╎   │   └─╼ 12
╎   └─╼ 8
╎   ├─╼ 11
╎   └─╼ 10
╙── 0
├─╼ 2
│   ├─╼ 6
│   └─╼ 5
└─╼ 1
├─╼ 4
└─╼ 3
╟── 0
╎   ├─╼ 1
╎   │   ├─╼ 3
╎   │   └─╼ 4
╎   └─╼ 2
╎   ├─╼ 5
╎   └─╼ 6
╙── 7
├─╼ 8
│   ├─╼ 10
│   └─╼ 11
└─╼ 9
├─╼ 12
└─╼ 13
"""
).strip()
assert ret == target

tree3 = nx.balanced_tree(r=2, h=2, create_using=nx.DiGraph)
tree3 = nx.relabel_nodes(tree3, {n: n + len(forest) for n in tree3.nodes})
forest = nx.union(forest, tree3)
ret = nx.forest_str(forest, sources=[0, 7, 14])
forest = nx.disjoint_union_all([tree1, tree2, tree3])
ret = nx.forest_str(forest, sources=[0, 14, 7])
print(ret)

target = dedent(
"""
╟── 0
╎   ├─╼ 1
╎   │   ├─╼ 3
╎   │   └─╼ 4
╎   └─╼ 2
╎   ├─╼ 5
╎   └─╼ 6
╟── 14
╎   ├─╼ 16
╎   │   ├─╼ 20
╎   │   └─╼ 19
╎   └─╼ 15
╎   ├─╼ 18
╎   └─╼ 17
╟── 7
╎   ├─╼ 9
╎   │   ├─╼ 13
╎   │   └─╼ 12
╎   └─╼ 8
╎   ├─╼ 11
╎   └─╼ 10
╙── 0
├─╼ 2
│   ├─╼ 6
│   └─╼ 5
└─╼ 1
├─╼ 4
└─╼ 3
╎   ├─╼ 15
╎   │   ├─╼ 17
╎   │   └─╼ 18
╎   └─╼ 16
╎   ├─╼ 19
╎   └─╼ 20
╙── 7
├─╼ 8
│   ├─╼ 10
│   └─╼ 11
└─╼ 9
├─╼ 12
└─╼ 13
"""
).strip()
assert ret == target
Expand All @@ -132,20 +130,20 @@ def test_undirected_multi_tree_forest():

target = dedent(
"""
╟── 7
╎   ├── 9
╎   │   ├── 13
╎   │   └── 12
╎   └── 8
╎   ├── 11
╎   └── 10
╙── 0
├── 2
│   ├── 6
│   └── 5
└── 1
├── 4
└── 3
╟── 0
╎   ├── 1
╎   │   ├── 3
╎   │   └── 4
╎   └── 2
╎   ├── 5
╎   └── 6
╙── 7
├── 8
│   ├── 10
│   └── 11
└── 9
├── 12
└── 13
"""
).strip()
assert ret == target
Expand All @@ -161,12 +159,12 @@ def test_undirected_tree_str():
node_target0 = dedent(
"""
╙── 0
├── 2
│   ├── 6
│   └── 5
└── 1
├── 4
└── 3
├── 1
│   ├── 3
│   └── 4
└── 2
├── 5
└── 6
"""
).strip()

Expand All @@ -179,12 +177,12 @@ def test_undirected_tree_str():
node_target2 = dedent(
"""
╙── 2
├── 6
├── 0
│   └── 1
│   ├── 3
│   └── 4
├── 5
└── 0
└── 1
├── 4
└── 3
└── 6
"""
).strip()
ret = nx.forest_str(graph, sources=[2])
Expand Down
44 changes: 24 additions & 20 deletions networkx/readwrite/text.py
Original file line number Diff line number Diff line change
Expand Up @@ -39,26 +39,27 @@ def forest_str(graph, use_labels=True, sources=None, write=None):
>>> graph = nx.balanced_tree(r=2, h=3, create_using=nx.DiGraph)
>>> print(nx.forest_str(graph))
╙── 0
├─╼ 2
│   ├─╼ 6
│   │   ├─╼ 14
│   │   └─╼ 13
│   └─╼ 5
│   ├─╼ 12
│   └─╼ 11
└─╼ 1
├─╼ 4
│   ├─╼ 10
│   └─╼ 9
└─╼ 3
├─╼ 8
└─╼ 7
├─╼ 1
│   ├─╼ 3
│   │   ├─╼ 7
│   │   └─╼ 8
│   └─╼ 4
│   ├─╼ 9
│   └─╼ 10
└─╼ 2
├─╼ 5
│   ├─╼ 11
│   └─╼ 12
└─╼ 6
├─╼ 13
└─╼ 14
>>> graph = nx.balanced_tree(r=1, h=2, create_using=nx.Graph)
>>> print(nx.forest_str(graph))
── 1
╎   ├── 2
╎   └── 0
── 0
── 1
└── 2
"""
import networkx as nx

Expand All @@ -83,11 +84,14 @@ def forest_str(graph, use_labels=True, sources=None, write=None):
sources = [n for n in graph.nodes if graph.in_degree[n] == 0]
else:
# use arbitrary sources for undirected trees
sources = sorted(graph.nodes, key=lambda n: graph.degree[n])
sources = [
min(cc, key=lambda n: graph.degree[n])
for cc in nx.connected_components(graph)
]

seen = set()
stack = []
for idx, node in enumerate(sources):
for idx, node in enumerate(sources[::-1]):
islast_next = idx == 0
stack.append((node, "", islast_next))

Expand Down Expand Up @@ -136,7 +140,7 @@ def forest_str(graph, use_labels=True, sources=None, write=None):
_write(this_prefix + str(label))

children = [child for child in succ[node] if child not in seen]
for idx, child in enumerate(children, start=1):
for idx, child in enumerate(children[::-1], start=1):
islast_next = idx <= 1
try_frame = (child, next_prefix, islast_next)
stack.append(try_frame)
Expand Down

0 comments on commit 16ac464

Please sign in to comment.