Recursive function for trees in Python

Question

I'm trying to make a function in Python, that takes an arbitrary node of a tree, and populates a list of lists based on the node give.

Given the following badly drawn tree:

If we start at, for example, node 5, we should get:

• A list containing all nodes with the same parent node, including the one we started at (4 and 5)
• Any child nodes, but not their children (6).
• Parent nodes and any parent nodes with the same parent, and their parent's nodes, etc. until we get to the root node, but not including the root node (just 2 and 3 in this case, but if the tree was deeper and we started lower, there'd be more here.

And the nodes should end up in a list of lists, one list for each depth.

The nodes in python:

nodes = [
    {'id': 1, 'parent': None},
    {'id': 2, 'parent': 1},
    {'id': 3, 'parent': 1},
    {'id': 4, 'parent': 2},
    {'id': 5, 'parent': 2},
    {'id': 6, 'parent': 5},
    {'id': 7, 'parent': 6},
    {'id': 8, 'parent': 3}
]

We can only see parents, not children, but this is the data format I have to work with, sadly.

So from this, if we take node 5, we want to end up with a node list looking something like this:

nl = [
    [{'id': 6, 'parent': 5}],
    [{'id': 4, 'parent': 2}, {'id': 5, 'parent': 2}],
    [{'id': 2, 'parent': 1}, {'id': 3, 'parent': 1}],
]

This is the code I have so far. I figure a recursive function is probably the simplest way. Unfortunately it seems to do nothing like what I think it should, and obviously I'm doing something very wrong. And this code doesn't even consider geting the child nodes which I'm not entirely sure how to deal with at all, apart from possibly handing the afterwards which would be much easier.

node_list = []

def pop_list(nodes=None, parent=None, node_list=None):
    if parent is None:
        return node_list
    node_list.append([])
    for node in nodes:
        if node['parent'] == parent:
            node_list[-1].append(node)
        if node['id'] == parent:
            parent = node['parent']
    return pop_list(nodes, parent, node_list)

print pop_list(nodes, 5, node_list)

Here is the output:

[[], [{'id': 3, 'parent': 1}], []]

Not exactly sure where I'm going wrong here.

Solution

The problem is here

    if node['id'] == parent:
        parent = node['parent']

The current parent will be overwritten by its parent.

Moreover, you should add return node_list at the end of the function, or use node_list as results.

def pop_list(nodes=None, parent=None, node_list=None):
    if parent is None:
        return node_list
    node_list.append([])
    for node in nodes:
        if node['parent'] == parent:
            node_list[-1].append(node)
        if node['id'] == parent:
            next_parent = node['parent']

    pop_list(nodes, next_parent, node_list)
    return node_list

>>> print pop_list(nodes, 5, node_list)
[[{'id': 6, 'parent': 5}], [{'id': 4, 'parent': 2}, {'id': 5, 'parent': 2}], [{'id': 2, 'parent': 1}, {'id': 3, 'parent': 1}]]  

Comments

[Tom Carrick]

Ah, a rookie mistake. Thanks!