python - 遍历和修改字典结构的树状列表

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
            yield from dfs(node['children'])

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

Python 3.3 及更高版本可以使用上面的递归 yield(yield from 语法)。早期版本会循环遍历递归结果，产生那些:

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

如果要选择的 ID 列表很大，将其从列表转换为集合会更高效，这将加快查找速度(所选中的 node['id'] ).

selected = set(selected)