题意
第一行是两个整数N和S,其中N是树的节点数。
第二行是N个正整数,第i个整数表示节点i的正整数。
接下来的N-1行每行是2个整数x和y,表示y是x的儿子。
输出格式:
输出路径节点总和为S的路径数量。
题解
因为求的路径都是从在包含根的一条链上,而且是从上往下。所以可以考虑直接dfs。记录元素前缀和。
比如搜到x,在set里面查询有没有s[x]-k,有的话ans++,然后继续往下dfs时插入s[x]。
最开始要插入0,同时在离开点x时要记住删除那个s[x]。
代码
#include<cstdio>
#include<cstring>
#include<algorithm>
#include<set>
#define N 100010
using namespace std;
multiset<int> q;
struct node{int to,next;}e[N*2];
int head[N],m;
int ans,a[N],n,s,sum[N];
void add_edge(int from,int to) {
e[++m].next = head[from];
head[from] = m;
e[m].to = to;
}
void dfs(int v) {
if(q.find(sum[v]-s) != q.end()) ans++;
q.insert(sum[v]);
for(int i = head[v];i;i=e[i].next) {
sum[e[i].to] = sum[v] + a[e[i].to];
dfs(e[i].to);
}
q.erase(q.find(sum[v]));
}
int main()
{
int u,v;
scanf("%d%d",&n,&s);
for(int i = 1;i <= n;i++) scanf("%d",&a[i]);
m = 0;
for(int i = 1;i < n;i++) {
scanf("%d%d",&u,&v);
add_edge(u,v);
}
q.insert(0); ans = 0;
sum[1] = a[1]; dfs(1);
printf("%d\n",ans);
return 0;
}