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

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

我喜欢在递归函数中使用闭包。对于这个例子来说，这并不是特别重要，但这样可以省去在递归调用时传递'选中的'这个参数。在更复杂的例子中，你可以在外层函数中保存很多状态，以便递归使用。

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

Answer 2

因为你是通过修改输入对象来操作的，而在Python中，对象是通过引用来处理的，所以你不需要在递归步骤中返回一个值或者使用返回值。此外，如果你能把子节点中的'None'替换成'[]'（更好的是，整个过程中都用元组而不是列表），那么你可以简化逻辑——这样你根本不需要一个基本情况，因为你可以递归到一个“空树”，这时它会对零个项目执行for循环，也就是啥都不做——这正是你想要的。

还有，为什么不使用Python内置的布尔类型呢？

def mark_selected(tree, selected):
    for obj in tree:
        obj['selected'] = obj['id'] in selected
        mark_selected(obj['children'], selected)

（哦，对了，你真的需要保持子节点的特定顺序吗？拥有一个包含'id'键的字典列表是不自然的；更合理的做法是用一个字典，其中键是id，值是没有'id'的字典。）

Answer 3

递归是一种非常优雅的编程方式。列表推导式不适合这里的情况，因为你是在原地改变结构，而不是生成一个新的序列。至于生成器，你可以写一个深度优先搜索（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 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'] in selected）。

selected = set(selected)