问题描述
我正在解决以下问题:
给定一棵三叉树(树的每个节点最多有三个 子),找到所有根到叶的路径。
示例:
我的代码如下:
from __future__ import annotations
import itertools
from collections import deque,Iterable
class TernaryNode:
def __init__(self,val: int) -> None:
self.children: list[TernaryNode] = []
self.val = val
def __repr__(self) -> str:
return str(self.val)
def ternary_tree_paths(root: TernaryNode) -> Iterable[Iterable[int]]:
def _visit(node: TernaryNode) -> Iterable[deque[int]]:
if not node:
return []
if not node.children:
queue = deque()
queue.append(node.val)
return [queue]
# **
paths = itertools.chain.from_iterable(map(lambda ch: _visit(ch),node.children))
for p in paths:
p.appendleft(node.val)
return paths
return _visit(root)
如图所示,上面的代码返回一个空列表,其中所需的行为是 [deque([1,2,3]),deque([1,4]),6])]
。注意带**的行;如果我将该行重写为 paths = [p for ch in node.children for p in _visit(ch)]
,它会按预期工作。我猜这个问题是因为函数 from_iterable
被懒惰地评估,但是当我迭代项目时不应该强制评估吗?
解决方法
在尝试对每个项目执行 appendleft
时,您正在耗尽可迭代链。之后 paths
为空。
您需要确保可迭代对象只计算一次:
def ternary_tree_paths(root: TernaryNode) -> Iterable[Iterable[int]]:
def _visit(node: TernaryNode) -> Iterable[deque[int]]:
if not node:
return []
if not node.children:
queue = deque()
queue.append(node.val)
return [queue]
paths = itertools.chain.from_iterable(map(_visit,node.children))
retval = [] # to keep track of results
for p in paths: # iterate
p.appendleft(node.val)
retval.append(p) # add to result
return retval # return result
return _visit(root)
这产生:
[deque([1,2,3]),deque([1,4]),6])]
对于问题中的示例。