在二叉树中找到累加和为指定值的最长路径长度

在二叉树中找到累加和为指定值的最长路径长度

题目描述

给定一颗二叉树和一个整数 sum，求累加和为 sum 的最长路径长度。路径是指从某个节点往下，每次最多选择一个孩子节点或者不选所形成的节点链。

输入描述：

第一行输入两个整数 n 和 root，n 表示二叉树的总节点个数，root 表示二叉树的根节点。
以下 n 行每行四个整数 fa，lch，rch，val，表示 fa 的左儿子为 lch，右儿子为 rch。val 表示 fa 节点的值(如果 lch 为 0 则表示 fa 没有左儿子，rch同理)

输出描述：

输出一个整数表示最长链的长度。

示例1

输入

9 1
1 2 3 -3
2 4 5 3
4 0 0 1
5 8 9 0
8 0 0 1
9 0 0 6
3 6 7 -9
6 0 0 2
7 0 0 1
6

输出

4

备注：

$1\le n\le 2\ast 10^5$

$1\le fa,lch,rch,root\le n$

$-10^6\le val\le 10^6$

---

题解：

如果一条树链（从根节点开始，到某个子节点结束）的累加和为 sum，而它的某条后缀树链累加和为 k ，那么前一半树链累加和为 sum - k（前缀和思想）。因此我们可以先序遍历二叉树，记录累加和，查看到当前节点时，是否出现过累加和为 sum-k 的前缀树链。

代码:

#include <cstdio>
#include <vector>
#include <unordered_map>

using namespace std;

const int N = 200010;

struct BST {
    int val, lch, rch;
} BST[N];
int n, rt, k, ret;
unordered_map<int, int> _hash;

void preorder(int root, int dep, int sum) {
    if (_hash.find(sum - k) != _hash.end()) ret = max(ret, dep - _hash[sum - k]);
    if (_hash.find(sum) == _hash.end()) _hash[sum] = dep;
    if (BST[root].lch) preorder(BST[root].lch, dep + 1, sum + BST[BST[root].lch].val);
    if (BST[root].rch) preorder(BST[root].rch, dep + 1, sum + BST[BST[root].rch].val);
    if (_hash[sum] == dep) _hash.erase(sum);
}

void solve() {
    scanf("%d%d", &n, &rt);
    int fa, lch, rch, val;
    while (n--) {
        scanf("%d %d %d %d", &fa, &lch, &rch, &val);
        BST[fa].lch = lch;
        BST[fa].rch = rch;
        BST[fa].val = val;
    }
    scanf("%d", &k);
    _hash[0] = 0;
    preorder(rt, 1, BST[rt].val);
    printf("%d\n", ret);
}

int main(void) {
    solve();
}