【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;
}