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

  |  

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

【对算法,数学,计算机感兴趣的同学,欢迎关注我哈,阅读更多原创文章】
我的网站:潮汐朝夕的生活实验室
我的公众号:算法题刷刷
我的知乎:潮汐朝夕
我的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