BZOJ 1316: 树上的询问 （点分治+set）

题目：http://www.lydsy.com/JudgeOnline/problem.php?id=1316

因为只要求存在某条路径长度为K，所以点分，然后用set判断差值是否在set中就可以了。

#include<cstring>
#include<algorithm>
#include<iostream>
#include<cstdio>
#include<queue>
#include<set>
#include<cmath>
#define rep(i,l,r) for (int i=l;i<=r;i++)
#define down(i,l,r) for (int i=l;i>=r;i--)
#define clr(x,y) memset(x,y,sizeof(x))
#define maxn 10900
#define inf int(1e9)
#define mm 1000000007
#define esp 1e-6
using namespace std;
#define ll long long
struct data{int obj,pre; ll c;
}e[maxn*2];
int head[maxn],s[maxn],ans[maxn],q[maxn],vis[maxn];
int n,m,tot,sum,mn,rt;
set<ll> bst;
int read(){
    int x=0,f=1; char ch=getchar();
    while (!isdigit(ch)) {if (ch=='-') f=-1; ch=getchar();}
    while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();
    return x*f;
}
void insert(int x,int y,int z){
    e[++tot].obj=y; e[tot].pre=head[x]; e[tot].c=z; head[x]=tot;
}
void dfs(int u,int fa){
    s[u]=1;  
    int mx=0;
    for (int j=head[u];j;j=e[j].pre){
        int v=e[j].obj;
        if (!vis[v]&&v!=fa){
            dfs(v,u);
            s[u]+=s[v];
             mx=max(mx,s[v]);
        }
    }
    mx=max(sum-mx,mx);
    if (mx<mn) mn=mx,rt=u;
}
void add(int u,int fa,ll w){
    bst.insert(w);
    for (int j=head[u];j;j=e[j].pre){
        int v=e[j].obj;
        if (v!=fa&&!vis[v])add(v,u,w+e[j].c);
    }
}
void update(int u,int fa,ll w){
    rep(i,1,m) if (!ans[i])
        ans[i]|=(bst.find(q[i]-w)!=bst.end());
    for (int j=head[u];j;j=e[j].pre){
        int v=e[j].obj;
        if (v!=fa&&!vis[v]) update(v,u,w+e[j].c);
    }
}
void solve(int u){
    mn=inf;
    dfs(u,0);
    u=rt;
    
    bst.clear(); bst.insert(0);
    for (int j=head[u];j;j=e[j].pre){
        int v=e[j].obj;
        if (!vis[v]) update(v,u,e[j].c),add(v,u,e[j].c);     //先判断点再加边。
    }
    
    vis[u]=1;
    for (int j=head[u];j;j=e[j].pre){
        int v=e[j].obj;
        if (!vis[v]){
            sum=s[v];
            solve(v);
        }
    }
    
}
int main(){
    n=read(); m=read();
    int x,y,z;
    rep(i,1,n-1){
        x=read(); y=read(); z=read();
        insert(x,y,z);
        insert(y,x,z);
    }
    rep(i,1,m) q[i]=read();
    sum=n;

    solve(1);
    rep(i,1,m) if(!q[i]) ans[i]=1;   //如果q[i]=0，那么可以只选取一个点。
    rep(i,1,m) if (ans[i]) puts("Yes"); else puts("No");
    return 0;
}