題目描述
在這個問題中,給定一個值 s 和一棵樹。在樹的每個節點有一個權值,第 i 個點的權值a_i,問有多少條路徑的節點權值總和爲 s。路徑中節點的深度必須是升序的。假設節點 1 是根節點,根的深度是 0,它的兒子節點的深度爲 1。路徑不必一定從根節點開始。
輸入格式
第一行有兩個整數 n 和 s,其中 n 是樹的節點數。
第二行有 n 個整數,第 i 個整數 a_i 表示節點 i 的權值。
接下來的 (n−1) 行,每行有兩個個整數 x 和 y,表示 y 是 x 的兒子。
輸出格式
輸出路徑節點總和爲 s 的路徑數量。
記錄根結點到每個點的路徑長度,放入set中,如果在遍歷一個點時,在set中找到sum[i]-s,說明在遍歷的這條鏈裏面有一個點到該點的距離爲 s ,ans++;
然後離開這條鏈的時候,記得清除這條鏈的sum[]值;
其實這種題目完全可以先想想在線性區間是怎麼做的,再放在樹形結構上往往可以解決;
代碼:
#include<bits/stdc++.h>
#define LL long long
#define pa pair<int,LL>
#define ls k<<1
#define rs k<<1|1
#define inf 0x3f3f3f3f
using namespace std;
const int N=100100;
const int M=400100;
const int mod=1e9;
int n,s,head[N],cnt;
int a[N],in[N],rt,ans,sum[N];
set<int>se;
struct Node{
int to,nex;
}edge[N];
void add(int p,int q){
edge[cnt].to=q;
edge[cnt].nex=head[p];
head[p]=cnt++;
}
void dfs(int p){
for(int i=head[p];~i;i=edge[i].nex){
int q=edge[i].to;
sum[q]=sum[p]+a[q];
se.insert(sum[q]);
if(se.count(sum[q]-s)) ans++;
dfs(q);
se.erase(sum[q]);
}
}
int main(){
memset(head,-1,sizeof(head));
scanf("%d%d",&n,&s);
for(int i=1;i<=n;i++) scanf("%d",&a[i]);
for(int i=1;i<n;i++){
int x,y;
scanf("%d%d",&x,&y);
add(x,y),in[y]++;
}
for(int i=1;i<=n;i++) if(!in[i]) rt=i;
sum[rt]=a[rt];
se.insert(a[rt]);
se.insert(0);
dfs(rt);
printf("%d\n",ans);
return 0;
}