算法数据结构

树形前缀和：在树的DFS过程中维护祖先链的和

字数统计: 1.1k字   |   阅读时长: 4分
, , ,

摘要: 在树结构上维护前缀和的问题

【对算法，数学，计算机感兴趣的同学，欢迎关注我哈，阅读更多原创文章】
我的网站：潮汐朝夕的生活实验室
我的公众号：算法题刷刷
我的知乎：潮汐朝夕
我的github：FennelDumplings
我的leetcode：FennelDumplings

在文章 面试好题连载 中，我提到了我认为的适合作为面试题的标准，这里复习一下：

对于适合面试的题目，以下几个要点中 (1)(2)(5)必须满足，(3)(4)至少满足一个：

(1) 考查基础的数据结构和主流算法，避开邪门算法和脑筋急转弯
(2) 暴力方法很容易想到
(3) 最好有两种以上的优化方式
(4) 最好有两三道递进的题可以放到一起
(5) 最终代码不长

今天我们看一个比较基础的题，leetcode第437题，【437. 路径总和 III】。要解决本题，需要先将问题拆解出一个子问题，也就是给定一列数 [a1, a2, …, ak]，问有多少子数组的和为 target，这是单独的一道题。因此 (4) 满足。本题的主要算法是前缀和，哈希表，均为基础算法，因此 (1) 满足。

本题是树形前缀或祖先链的经典问题，在 DFS 的过程中维护祖先链上的和。下面我们来看一下这个题。

题目

给定一个二叉树的根节点 root ，和一个整数 targetSum ，求该二叉树里节点值之和等于 targetSum 的 路径 的数目。

路径 不需要从根节点开始，也不需要在叶子节点结束，但是路径方向必须是向下的（只能从父节点到子节点）。

1
2
3
二叉树的节点个数的范围是 [0,1000]
-1e9 <= Node.val <= 1e9 
-1000 <= targetSum <= 1000

示例 1：

输入：root = [10,5,-3,3,2,null,11,3,-2,null,1], targetSum = 8
输出：3
解释：和等于 8 的路径有 3 条，如图所示。
示例 2：
输入：root = [5,4,8,11,null,13,4,7,2,null,null,5,1], targetSum = 22
输出：3

题解

算法：前缀和 + 哈希映射

dfs(前序遍历)时，栈里存的是当前节点到根的链，这条链上所有元素的和可以作为前缀和维护在哈希映射 mapping 里，其中 mapping 的键是前缀和，值是出现的次数。

假设当前访问到节点 Ni，栈里存的从当前节点到根节点的链为 [N0, N1, N2, …, Ni-1, Ni]，此时在 mapping 中保存着前面 N0 到 Ni-1 的所有前缀和的次数信息。

此时 N0 到 Ni 的前缀和为 presum，查询 mapping 中有多少个 presum - target，这个个数就是以 Ni 结尾的和为 target 的链的个数。

查询完成后，将 presum 插入到 mapping 中，然后进入子树。

当 Ni 对应的子树都访问完，准备返回的时候，节点对应的前缀和 presum 要先从 unordered_map 中删掉。

算法中主要涉及到的算法要点是前缀和，文章 【模板】前缀和与差分 中总结了算法的原理和 C++/Python 的可以直接用的代码模板。本题中需要用哈希表将前缀和保存起来，属于【数据结构维护前缀和】的套路。前缀和在 leetcode 中有 50 道题左右，而且周赛经常会考，文章 前缀和问题分类汇总 中总结了前缀和的几乎所有套路，可以集中练习。

代码(C++)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
class Solution {
public:
    int pathSum(TreeNode* root, int targetSum) {
        if(!root) return 0;
        mapping = unordered_map<int, int>();
        mapping[0] = 1;
        int ans = 0;
        _preOrder(root, targetSum, 0, ans);
        return ans;
    }

private:
    unordered_map<int, int> mapping;

    void _preOrder(TreeNode* node, const int targetSum, int presum, int& ans)
    {
        presum += node -> val;
        ans += mapping[presum - targetSum];
        ++mapping[presum];
        if(node -> left)
            _preOrder(node -> left, targetSum, presum, ans);
        if(node -> right)
            _preOrder(node -> right, targetSum, presum, ans);
        --mapping[presum];
    }
};

代码(Python)

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
from collections import Counter

class Solution:
    def pathSum(self, root: TreeNode, targetSum: int) -> int:
        if root is None:
            return 0
        self.ans = 0
        self.mapping = Counter()
        self.mapping.update([0])
        self._preOrder(root, targetSum, 0)
        return self.ans

    def _preOrder(self, node: TreeNode, targetSum: int, presum: int) -> None:
        presum += node.val
        self.ans += self.mapping[presum - targetSum]
        self.mapping.update([presum])
        if node.left is not None:
            self._preOrder(node.left, targetSum, presum)
        if node.right is not None:
            self._preOrder(node.right, targetSum, presum)
        self.mapping.subtract([presum])
Share