Tree interval representation in python

Question

Consider I have a tree built with

import networkx as nx
dg = nx.DiGraph()
dg.add_edges_from([('a', 'b'), ('a', 'c'), ('b', 'd'), ('b', 'e')])

I want an interval representation of the following tree, what I mean is

[1, [2, [3, 4], [5, 6], 7], [8,9 ], 10]

Where the nesting corresponds to the tree.

Are there any functions in networkx or in an other library that allow to do this?

Answer 1

I don't think you will find a lib which do that. However it's easy with a DFS. Here is a recursive implementation:

def tree(dg):
    def walk(root, start):
        children = dg.successors(root)
        if children:
            cur = start+1
            res = [start]
            for c in children:
                l = walk(c, cur)
                res.append(l)
                cur = l[-1]+1
            res.append(cur)
            return res
        else:
            return [start, start+1]

    roots = [i for i in dg.nodes() if dg.in_degree(i) == 0]
    assert(len(roots) == 1)
    root = roots[0]

    return walk(root, 1)

Then:

>>> tree(dg)
[1, [2, 3], [4, [5, 6], [7, 8], 9], 10]

By the way, you did notice that my answer is not the same as yours. You didn't specify any ordering on the children. If you want to, just sort the children list according to the desired ordering.