Traversing and modifying a tree-like list of dict structure

Question

I have a structure that looks like this:

[ {'id': 4, 'children': None},
  {'id': 2, 'children': 
    [ {'id': 1, 'children':
        [ {'id': 6, 'children': None},
          {'id': 5, 'children': None} ]
      },
      {'id': 7, 'children':
        [ {'id': 3, 'children': None} ]
      }
    ]
  }
]

I also have a list of selected IDs, [4, 5, 6, 7]. I want to traverse the list and for each object in the list, add a selected key with a value of 1 if it is selected, and 0 if it is not.

Currently I am doing this recursively with this function:

def mark_selected(tree, selected):
    for obj in tree:
        obj['selected'] = 1 if obj['id'] in selected else 0
        if obj['children'] is not None:
            obj['children'] = mark_selected(obj['children'], selected)
    return tree

This seems to work fine, but I was wondering if there was a more clever way to do this, possibly using list comprehension or generators.

Can anyone come up with a more elegant solution for this?

Answer 1

Recursion is perfectly elegant. List comprehensions don't apply as you're altering the structure in place, rather than producing a new sequence. As for generators, you could write a DFS or BFS traverser.

def dfs(nodes):
    if nodes is not None:
        for node in nodes:
            yield node
            yield from dfs(node['children'])

for node in dfs(tree):
    node['selected'] = node['id'] in selected

Python 3.3 and later can use the recursive yield above (yield from syntax). Earlier versions would instead loop over the recursive results, yielding those:

def dfs(nodes):
    if nodes is not None:
        for node in nodes:
            yield node
            for child in dfs(node['children']):
                yield child

If the list of IDs to select is large, it would be more performant to convert it from a list to a set, which will speed up lookup (the node['id'] in selected).

selected = set(selected)