Search code examples
pythonnetworkxdigraphs

DiGraph parallel ordering

I have this kind of Directed Acyclic Graph with multiple roots:

DiGraph

And I need to get a list with nodes sorted by directions and grouped by steps, like this:

ordering = [
    [1, 3],
    [2],
    [4],
    [5, 6],
    [7]
]

Maybe there is some ready algorithm for this? I tried networkx.algorithm but all of them can return me only a flat list without grouping by steps.

Solution

  • nx.topological_sort almost does what you want; the only difference is that it doesn't group items that enter the queue at the same time, but it's straightforward to adapt the source to make it do so:

    def topological_sort_grouped(G):
    indegree_map = {v: d for v, d in G.in_degree() if d > 0}
    zero_indegree = [v for v, d in G.in_degree() if d == 0]
    while zero_indegree:
        yield zero_indegree
        new_zero_indegree = []
        for v in zero_indegree:
            for _, child in G.edges(v):
                indegree_map[child] -= 1
                if not indegree_map[child]:
                    new_zero_indegree.append(child)
        zero_indegree = new_zero_indegree

    With your example:

    In [21]: list(nx.topological_sort(G))
Out[21]: [3, 1, 2, 4, 6, 7, 5]

In [22]: list(topological_sort_grouped(G))
Out[22]: [[1, 3], [2], [4], [5, 6], [7]]

    In practice, I have to wonder if there's a situation where this construction is more useful than just using nx.topological_sort (or nx.lexicographical_topological_sort) directly?