Skip to content

Commit

Permalink
Add random_ordered_tree and forest_str
Browse files Browse the repository at this point in the history
  • Loading branch information
Erotemic committed Dec 3, 2020
1 parent ed8aee7 commit e518a7d
Show file tree
Hide file tree
Showing 5 changed files with 525 additions and 4 deletions.
45 changes: 45 additions & 0 deletions networkx/generators/random_graphs.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
"random_powerlaw_tree",
"random_powerlaw_tree_sequence",
"random_kernel_graph",
"random_ordered_tree",
]


Expand Down Expand Up @@ -1293,3 +1294,47 @@ def my_function(b):
j = int(math.ceil(n * kernel_root(i / n, j / n, r)))
graph.add_edge(i - 1, j - 1)
return graph


@py_random_state(2)
def random_ordered_tree(n, seed=None, directed=False):
"""
Creates a random ordered tree
Parameters
----------
n : int
A positive integer representing the number of nodes in the tree.
seed : integer, random_state, or None (default)
Indicator of random number generation state.
See :ref:`Randomness<randomness>`.
directed : bool
if the edges are one-way
Returns
-------
networkx.OrderedDiGraph | networkx.OrderedGraph
Example
-------
>>> import networkx as nx
>>> assert len(nx.random_ordered_tree(n=1, seed=0).nodes) == 1
>>> assert len(nx.random_ordered_tree(n=2, seed=0).nodes) == 2
>>> assert len(nx.random_ordered_tree(n=3, seed=0).nodes) == 3
>>> otree = nx.random_ordered_tree(n=5, seed=3, directed=True)
>>> print(nx.forest_str(otree))
╙── 0
└─╼ 1
└─╼ 4
├─╼ 2
└─╼ 3
"""
from networkx.utils import create_py_random_state

rng = create_py_random_state(seed)
# Create a random undirected tree
create_using = nx.OrderedDiGraph if directed else nx.OrderedGraph
otree = nx.random_tree(n, seed=rng, create_using=create_using)
return otree
66 changes: 62 additions & 4 deletions networkx/generators/trees.py
Original file line number Diff line number Diff line change
Expand Up @@ -153,7 +153,7 @@ def _helper(paths, root, B):
# > method of generating uniformly distributed random labelled trees.
#
@py_random_state(1)
def random_tree(n, seed=None):
def random_tree(n, seed=None, create_using=None):
"""Returns a uniformly random tree on `n` nodes.
Parameters
Expand Down Expand Up @@ -184,11 +184,69 @@ def random_tree(n, seed=None):
*n* nodes, the tree is chosen uniformly at random from the set of
all trees on *n* nodes.
Example
-------
>>> import networkx as nx
>>> tree = nx.random_tree(n=10, seed=0)
>>> print(nx.forest_str(tree, sources=[0]))
╙── 0
├── 3
└── 4
├── 6
│   ├── 1
│   ├── 2
│   └── 7
│   └── 8
│   └── 5
└── 9
>>> import networkx as nx
>>> tree = nx.random_tree(n=10, seed=0, create_using=nx.OrderedDiGraph)
>>> print(nx.forest_str(tree))
╙── 0
├─╼ 3
└─╼ 4
├─╼ 6
│   ├─╼ 1
│   ├─╼ 2
│   └─╼ 7
│   └─╼ 8
│   └─╼ 5
└─╼ 9
"""
if n == 0:
raise nx.NetworkXPointlessConcept("the null graph is not a tree")
# Cannot create a Prüfer sequence unless `n` is at least two.
if n == 1:
return nx.empty_graph(1)
sequence = [seed.choice(range(n)) for i in range(n - 2)]
return nx.from_prufer_sequence(sequence)
utree = nx.empty_graph(1)
else:
sequence = [seed.choice(range(n)) for i in range(n - 2)]
utree = nx.from_prufer_sequence(sequence)

if create_using is None:
tree = utree
else:
# TODO: maybe a tree classmethod like
# Graph.new, Graph.fresh, or something like that
def new(cls_or_self):
if hasattr(cls_or_self, "_adj"):
# create_using is a NetworkX style Graph
cls_or_self.clear()
self = cls_or_self
else:
# try create_using as constructor
self = cls_or_self()
return self

tree = new(create_using)
if tree.is_directed():
# Use a arbitrary root node and dfs to define edge directions
edges = nx.dfs_edges(utree, source=0)
else:
edges = utree.edges

# Populate the specified graph type
tree.add_nodes_from(utree.nodes)
tree.add_edges_from(edges)

return tree
1 change: 1 addition & 0 deletions networkx/readwrite/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,3 +16,4 @@
from networkx.readwrite.gexf import *
from networkx.readwrite.nx_shp import *
from networkx.readwrite.json_graph import *
from networkx.readwrite.text import *

0 comments on commit e518a7d

Please sign in to comment.