遍历和修改类似于dict结构的树状列表

Question

我有一个看起来像这样的结构:

[ {'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} ]
      }
    ]
  }
]

我还有一个选定的ID列表[4, 5, 6, 7].我想遍历列表,并为列表中的每个对象添加selected一个值,1如果选中它,0如果不是,则添加值.

目前我正在使用此函数递归执行此操作:

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

这似乎工作正常,但我想知道是否有更聪明的方法来做到这一点,可能使用列表理解或生成器.

任何人都可以为此提出更优雅的解决方案吗？

Answer 1

递归非常优雅.列表推导不适用,因为您正在改变结构,而不是生成新的序列.对于生成器,您可以编写DFS或BFS遍历器.

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

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

如果要选择的ID列表很大,那么将ID转换为具有ID作为键的dict将更加高效,这将加速查找(the node['id'] in selected).

selected = dict(zip(selected, selected))