LeetCode-python 面试题 04.12. 求和路径

题目链接
难度：中等       类型： DFS, 前缀和

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

示例：

解题思路

最朴素的思路以每一个结点为开始，遍历所有路径，这会有很多重复计算，当前节点到根节点的和可以记录下来，即前缀和

代码实现

基本的dfs解法

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> int:
         
        def dfs(node, cur_sum):
            cnt = 0
            if node is None:
                return 0
            cur_sum += node.val
            if cur_sum == sum:
                cnt += 1

            cnt += dfs(node.left, cur_sum)
            cnt += dfs(node.right, cur_sum)
            return cnt

        if root is None:
            return 0
      
        cnt = dfs(root, 0)

        cnt += self.pathSum(root.left, sum)
        cnt += self.pathSum(root.right, sum)
     
        return cnt

加入前缀和的dfs

class Solution:
    def pathSum(self, root: TreeNode, sum: int) -> int:
        prefix_sum = collections.defaultdict(int)
        prefix_sum[0] = 1
        def dfs(node, cur_sum):
            cnt = 0
            if node is None:
                return 0
            cur_sum += node.val
            cnt += prefix_sum[cur_sum - sum]
            prefix_sum[cur_sum] += 1
            cnt += dfs(node.left, cur_sum)
            cnt += dfs(node.right, cur_sum)
            # 唯一的关键点在于，退出当前结点时，需要删除当前节点的记录
            prefix_sum[cur_sum] -= 1
            return cnt

        if root is None:
            return 0
 
        return dfs(root, 0)