樹形結構+前綴和 [JLOI2012]樹(洛谷 P3252)

題目描述

在這個問題中,給定一個值 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;
}
發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章