[JLOI2012] 树

Description

在这个问题中，给定一个值S和一棵树。在树的每个节点有一个正整数，问有多少条路径的节点总和达到S。路径中节点的深度必须是升序的。假设节点1是根节点，根的深度是0，它的儿子节点的深度为1。路径不必一定从根节点开始。

Input

第一行是两个整数N和S，其中N是树的节点数。 第二行是N个正整数，第i个整数表示节点i的正整数。 接下来的N-1行每行是2个整数x和y，表示y是x的儿子。

Output

输出路径节点总和为S的路径数量。

Range

对于100%数据，N<=100000，所有权值以及S都不超过1000。

Solution

转化一下题意，就是求树上的一条链，使权值之和等于s。

我们可以利用dfs求出树上每个点到根的权值和(也就是树上的前缀和)，在回溯的过程中求出答案即可。

具体做法是，当我们在搜一个点 i 时，看一眼有没有它的某个祖先 j 使得 qzh[j]-qzh[i]=p

如何找这个祖先呢？我们可以用 STL 中的 set ，在 dfs 的时候将当前点的前缀和插进集合，回溯的时候找集合中是否有值为 p-qzh[now] 的点，如果有，代表它的某个祖先到它即为一条合法路径， ans++，回溯最后记得从集合中 erase 掉 qzh[now] 即可。

Code

#include<set>
#include<cstdio>
#define N 100005
#define int long long
using namespace std;

int ans;
set<int> s;
int head[N];
int is_root[N];
int val[N],qzh[N];
;

struct Edge{
    int to,nxt;
}edge[];

void add(int x,int y){
    edge[++cnt].to=y;
    edge[cnt].nxt=head[x];
    head[x]=cnt;
}

void dfs(int now,int fa){
    for(int i=head[now];i;i=edge[i].nxt){
        if(edge[i].to==fa) continue;
        s.insert(qzh[now]);
        qzh[edge[i].to]=qzh[now]+val[edge[i].to];
        dfs(edge[i].to,now);
        int k=qzh[edge[i].to]-p;
        if(s.count(k)) ans++;
        s.erase(qzh[edge[i].to]);
    }
}

signed main(){
    scanf("%lld%lld",&n,&p);
    ;i<=n;i++) scanf("%lld",&val[i]);
    ;i<n;i++){
        scanf("%lld%lld",&x,&y);
        add(x,y);add(y,x);
    }
    s.insert();
    qzh[]=val[];
    dfs(,);
    printf("%lld",ans);
    ;
}