Search code examples
pythonlistrecursiontree

Unravel a tree structure stored as nested list


I'm working on a piece of code which simulates possible paths through a series of nodes. The nodes are assigned numbers and the connected nodes are listed in a lookup table. Since the path sometimes splits in two (or many) possibilities (and sometimes only one, i.e., doesn't split) I get a tree-type structure in the form of nested lists. For each level in the list I get a list of possible continuations. For example, a tree of paths pay look like L = [70, [49, [4]], [94, [120, [144]],[121, [150,[173]]]]], i.e., consists of three individual paths [70,49,4], [70,94,120,144], and [70,94,121,150,173].

Is there a nice way to unravel this tree so I can get back the individual paths? The returned format is not very important, as long as I can use it to make a visualization of the nodes changing colors. Note that I have no idea of how long the paths (the depth of the list) may be, so recursion is probably a good idea. Also I don't know how many nodes there are in any one path or how many (allowed) split roads a node will have, except that all those numbers are finite. One way to obtain it manually for the second branch is L[0], L[2][0], L[2][1][0],L[2][1][1][0], if that can be of any help.


Solution

  • You could indeed use recursion for this, for instance with a recursive generator:

    def unravel(lst):
        if len(lst) > 1:
            for child in lst[1:]:
                for path in unravel(child):
                    yield [lst[0], *path]
        else:
            yield lst
    

    Call like this:

    lst = [70, [49, [4]], [94, [120, [144]], [121, [150,[173]]]]]
    
    for path in unravel(lst):
        print(path)