三叉树路径

问题描述

我正在解决以下问题:

给定一棵三叉树(树的每个节点最多有三个 子),找到所有根到叶的路径。

示例:

enter image description here

我的代码如下:

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])]

对于问题中的示例。