1"""Functions for generating stochastic graphs from a given weighted directed
2graph.
3
4"""
5
6from networkx.classes import DiGraph
7from networkx.classes import MultiDiGraph
8from networkx.utils import not_implemented_for
9
10__all__ = ["stochastic_graph"]
11
12
13@not_implemented_for("undirected")
14def stochastic_graph(G, copy=True, weight="weight"):
15    """Returns a right-stochastic representation of directed graph `G`.
16
17    A right-stochastic graph is a weighted digraph in which for each
18    node, the sum of the weights of all the out-edges of that node is
19    1. If the graph is already weighted (for example, via a 'weight'
20    edge attribute), the reweighting takes that into account.
21
22    Parameters
23    ----------
24    G : directed graph
25        A :class:`~networkx.DiGraph` or :class:`~networkx.MultiDiGraph`.
26
27    copy : boolean, optional
28        If this is True, then this function returns a new graph with
29        the stochastic reweighting. Otherwise, the original graph is
30        modified in-place (and also returned, for convenience).
31
32    weight : edge attribute key (optional, default='weight')
33        Edge attribute key used for reading the existing weight and
34        setting the new weight.  If no attribute with this key is found
35        for an edge, then the edge weight is assumed to be 1. If an edge
36        has a weight, it must be a positive number.
37
38    """
39    if copy:
40        G = MultiDiGraph(G) if G.is_multigraph() else DiGraph(G)
41    # There is a tradeoff here: the dictionary of node degrees may
42    # require a lot of memory, whereas making a call to `G.out_degree`
43    # inside the loop may be costly in computation time.
44    degree = dict(G.out_degree(weight=weight))
45    for u, v, d in G.edges(data=True):
46        if degree[u] == 0:
47            d[weight] = 0
48        else:
49            d[weight] = d.get(weight, 1) / degree[u]
50    return G
51