求和路径

标签: 深度优先搜索 二叉树

难度: Medium

给定一棵二叉树,其中每个节点都含有一个整数数值(该值或正或负)。设计一个算法,打印节点数值总和等于某个给定值的所有路径的数量。注意,路径不一定非得从二叉树的根节点或叶节点开始或结束,但是其方向必须向下(只能从父节点指向子节点方向)。

示例:
给定如下二叉树,以及目标和 sum = 22

              5
             / \
            4   8
           /   / \
          11  13  4
         /  \    / \
        7    2  5   1

返回:

3
解释:和为 22 的路径有:[5,4,11,2], [5,8,4,5], [4,11,7]

提示:

  • 节点总数 <= 10000

Submission

运行时间: 30 ms

内存: 17.0 MB

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> int:
        self.ans=0
        prefix=defaultdict(int)
        prefix[0]=1
        def path(root, total):
            if not root:
                return
            total+=root.val
            self.ans+=prefix[total-sum]
            prefix[total]+=1
            path(root.left, total)
            path(root.right, total)
            prefix[total]-=1
        
        path(root, 0)
        return self.ans

Explain

该题解采用了前缀和的方法来高效地计算路径和。对于二叉树中的任意一条路径,可以通过维护一个从根节点到当前节点的路径和,以及一个哈希表来记录所有可能的前缀和及其出现的次数。对每个节点,计算当前的路径和,并查询哈希表中路径和减去目标值(sum)的结果出现的次数,这些次数即为以当前节点结束的、和为目标值的路径数量。在遍历每个节点后,更新哈希表以包含当前路径和的次数,并在返回前减少当前路径和的计数,以防影响其他路径的计算。

时间复杂度: O(n)

空间复杂度: O(n)

# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> int:
        self.ans=0  # 用于存储路径和等于目标值的路径数量
        prefix=defaultdict(int)  # 哈希表用于记录前缀和及其出现的次数
        prefix[0]=1  # 初始化前缀和为0的情况,对应从根节点开始的路径
        def path(root, total):
            if not root:  # 如果节点为空,则返回
                return
            total+=root.val  # 更新当前路径和
            self.ans+=prefix[total-sum]  # 查找当前路径和减去目标值的结果在哈希表中的次数
            prefix[total]+=1  # 更新哈希表,包含当前路径和的次数
            path(root.left, total)  # 遍历左子树
            path(root.right, total)  # 遍历右子树
            prefix[total]-=1  # 回溯,减少当前路径和的计数
        
        path(root, 0)  # 从根节点开始遍历
        return self.ans  # 返回结果

Explore

在哈希表中初始设置prefix[0] = 1是为了处理从根节点开始的路径正好等于目标和的情况。这个初始化意味着存在一个虚拟的前缀和为0,它代表在任何实际节点之前,没有任何节点的和。因此,当某个节点的累计和等于目标和时,它可以被视为从根节点到该节点的路径和等于目标和,而这条路径的前缀和与'0'的差也恰好等于目标和。

在函数path中,使用prefix[total-sum]的查询确保找到的路径是向下且不跨越非父子节点的,这是因为前缀和的更新和查询都是在单个从根到叶的递归遍历过程中进行的。只有当前遍历的路径的节点会影响当前的路径和total,并且prefix哈希表是实时更新的,每次递归返回前会将当前节点的路径和计数减少,确保不会影响其他分支的计算。这样,每次查询prefix[total - sum]时,它反映的是当前节点的祖先节点到当前节点的路径和,且这些路径都是从上到下的。

在递归函数path中,对prefix的更新操作在递归调用前后都执行,这是为了正确管理路径和的计数,以便能够准确地计算符合条件的路径数量。在递归调用之前,需要增加当前路径和的计数,这是因为当前节点已被纳入路径和中。而递归调用之后,需要减少当前路径和的计数是为了进行回溯,这确保了当前节点的路径和不会影响到其它可能的路径计算。这种'增加-递归-减少'的模式是为了保证每个节点在递归搜索过程中只影响它所在路径的计算,不会错误地影响其他路径。