【bzoj1316】【树上的询问】点分治+set

【bzoj1316】【树上的询问】点分治+set

Description
一棵n个点的带权有根树，有p个询问，每次询问树中是否存在一条长度为Len的路径，如果是，输出Yes否输出No.
Input
第一行两个整数n, p分别表示点的个数和询问的个数． 接下来n-1行每行三个数x, y, c，表示有一条树边x→y，长度为c． 接下来p行每行一个数Len，表示询问树中是否存在一条长度为Len的路径．
Output
输出有p行，Yes或No.
Sample Input
6 4 
1 2 5 
1 3 7 
1 4 1 
3 5 2 
3 6 3 
1 
8 
13 
14 
Sample Output
Yes 
Yes 
No 
Yes 

HINT
30%的数据，n≤100． 
100%的数据，n≤10000，p≤100，长度≤1000000． 

#include<bits/stdc++.h>
using namespace std;
struct Node{int to,v;};
set<int> res;
vector<Node> G[10005];
Node e;
int v[10005],n,m,sim[10005],mos[10005],MX=0x3f3f3f3f,root,a,b,w,dis[10005];
void getroot(int u,int fa){
	sim[u]=1;mos[u]=0;
	v[u]=1;
	for(int i=0;i<G[u].size();i++){
		e=G[u][i];
		if(v[e.to]||e.to==fa)continue;
		getroot(e.to,u);
		sim[u]+=sim[e.to];
		mos[u]=max(mos[u],sim[e.to]);
	}
	mos[u]=max(mos[u],n-sim[u]);
	if(MX>mos[u]){
		root=MX;
		MX=mos[u];
	}
}//找到一个根 
void getans(int u,int fa,int d,set<int> ans){
	for(int i=0;i<G[u].size();i++){
		e=G[u][i];
		if(e.to==fa)continue;
		ans.insert(d+e.v);
		getans(e.to,u,d+e.v,ans); 
	}	 
}
//首先对根进行加乘  再依次对每个点进行加乘 
void divide(int u,int fa){
	v[u]=1;
	vector<set<int> > all; 
	for(int i=0;i<G[u].size();i++){
		set<int> ans;
		e=G[u][i];
		if(v[e.to]||e.to==fa)continue;
		ans.insert(e.v);//每条边都加上
		getans(e.to,u,e.v,ans);//找这个子树的所有通过节点u的边 
		all.push_back(ans);//子树所有点距离该点的长度加到vector里面 
	}
	cout<<all.size()<<endl; 
	for(int i=0;i<all.size();i++){
		set<int> ans=all[i];
		for(set<int>::iterator it=ans.begin();it!=ans.end();it++){
			for(int j=i+1;j<all.size();j++){
				for(set<int>::iterator ti=(all[j]).begin();ti!=(all[j]).end();ti++){
					res.insert(*it+*ti);
				}
			}
		}
	}
}

int main(){
	cin>>n>>m;
	for(int i=1;i<n;i++){
		cin>>a>>b>>w;
		e.to=b;e.v=w;
		G[a].push_back(e);
		e.to=a;
		G[b].push_back(e);
	}
	getroot(1,0);
	memset(v,0,sizeof(v));
	divide(root,0);//初始化完毕，开始进行输出 
	cout<<res.size()<<endl;
	for(int i=1;i<=m;i++){
		cin>>a;
		if(res.count(a)) cout<<"Yes"<<endl;
		else cout<<"No"<<endl;
	}
	return 0;
}