门|门 IT 2005 |问题 10

问题描述

在一个由若干个节点组成的二叉树中，每个节点只有一个数字，已知该二叉树的根节点，求二叉树中所有从根节点到叶子节点的路径上数字的和等于给定值的路径。

请你编写一个函数 solution，它接收以下参数：

• root: 表示二叉树的根节点。
• targetSum: 表示要求的路径上数字的和。

该函数需要返回二叉树中所有满足要求的路径，每条路径需按照节点的顺序组成一个列表。如果不存在满足要求的路径，返回一个空列表。

函数签名

from typing import List

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def solution(root: TreeNode, targetSum: int) -> List[List[int]]:
    pass

示例

root = TreeNode(5)

root.left = TreeNode(4)
root.right = TreeNode(8)

root.left.left = TreeNode(11)
root.right.left = TreeNode(13)
root.right.right = TreeNode(4)

root.left.left.left = TreeNode(7)
root.left.left.right = TreeNode(2)
root.right.right.left = TreeNode(5)
root.right.right.right = TreeNode(1)

assert solution(root, 22) == [
    [5, 4, 11, 2],
    [5, 8, 4, 5]
]
assert solution(root, 14) == [[5, 4, 11]]
assert solution(root, 27) == [[5, 8, 4, 5, 1]]
assert solution(root, 100) == []

解题思路

可以使用深度优先搜索来遍历整个二叉树，同时记录下从根节点到当前节点的路径及其路径上所有节点的值的总和。

具体来说，我们可以维护一个栈，其中每个元素都包含当前节点以及从根节点到当前节点的路径和。根据栈的后进先出的特点，深度优先搜索时，每次取出栈顶的元素，并拓展出该节点的所有子节点到栈中。如果当前节点是叶子节点并且路径和等于目标值，则找到了一条满足条件的路径；如果当前节点不是叶子节点，则将其子节点加入栈顶，并更新子节点的路径和。

代码如下：

def solution(root: TreeNode, targetSum: int) -> List[List[int]]:
    if not root:
        return []

    res = []
    stack = [(root, [root.val])]

    while stack:
        node, path = stack.pop()

        if not node.left and not node.right:  # 叶子节点
            if sum(path) == targetSum:
                res.append(path)
        else:
            if node.left:
                stack.append((node.left, path + [node.left.val]))
            if node.right:
                stack.append((node.right, path + [node.right.val]))

    return res

时间复杂度：$O(n \log n)$。

参考资料

• LeetCode 题解 113. 路径总和 II