Parse a tree & return the various paths

Question

Input:

{ 1: [2,3], 2:[4], 3: [5,6], 4:[6] }

Expected output:

[[1, 2, 4, 6], [1, 3, 5], [1, 3, 6]]

My Code:

def parse(node, tree):
    result = []
    if node not in tree:
        return [[node]]
    else:
        res = []
        for next_node in tree[node]:
            res.extend(parse(next_node, tree))
        for r in res:
           result.append([node]+r)
        return result

I'm looking for a less-complex solution.

Answer 1

You can get rid of the variables result and res and shorten the code by changing it into a generator. If you really need to get a list in the end, wrap the call with list(parse(root, tree)).

def parse(node, tree):
    if node not in tree:
        yield [node]
    else:
        for next_node in tree[node]:
            for r in parse(next_node, tree):
                yield [node] + r

Answer 2

What exactly do you expect? The code is fairly concise and clear. It's complexity is also well. If anything, the list concatenation is what's taking the most time here. So I would just change it slightly to avoid that like this:

def parse(node, tree, depth=1):
    result = []
    if node not in tree:
        return [[node] * depth]
    else:
        res = []
        for next_node in tree[node]:
            res.extend(parse(next_node, tree, depth+1))
        for r in res:
           r[depth-1] = node
           result.append(r)
        return result

This way you automatically create the list of the correct length and avoid adding to a list over and over.