【BZOJ1316】—树上的询问（点分治）

传送门

对于每个根，我们找dis为len的点就可以了

先dfs一次，然后按距离排序后将所有为len的点的贡献加就可以了、

#include<bits/stdc++.h>
using namespace std;
#define ll long long
inline int read(){
    char ch=getchar();
    int res=0,f=1;
    while(!isdigit(ch)) {if(ch=='-') f=-1;ch=getchar();}
    while(isdigit(ch)) res=(res<<3)+(res<<1)+(ch^48),ch=getchar();
    return res*f;
}
const int N=10005;
int adj[N],nxt[N<<1],to[N<<1],son[N],siz[N],p,n,root,maxn,d[N],val[N<<1],dis[N],ans[105],q[105],cnt,ecnt;
bool vis[N];
inline void addedge(int u,int v,int w){
    nxt[++cnt]=adj[u],adj[u]=cnt,to[cnt]=v,val[cnt]=w;
    nxt[++cnt]=adj[v],adj[v]=cnt,to[cnt]=u,val[cnt]=w;
}
inline void getroot(int u,int fa){
    siz[u]=1,son[u]=0;
    for(int e=adj[u];e;e=nxt[e]){
        int v=to[e];
        if(vis[v]||v==fa) continue;
        getroot(v,u);
        siz[u]+=siz[v];
        son[u]=max(son[u],siz[v]);
    }
    son[u]=max(son[u],maxn-siz[u]);
    if(son[u]<son[root]) root=u;
}
inline void getdis(int u,int fa){
    d[++ecnt]=dis[u];
    for(int e=adj[u];e;e=nxt[e]){
        int v=to[e];
        if(vis[v]||v==fa) continue;
        dis[v]=dis[u]+val[e];
        getdis(v,u);
    }
}
inline int findl(int l,int r,int k){
    int res=0;
    while(l<=r){
        int mid=(l+r)/2;
        if(d[mid]==k){
            res=mid,r=mid-1;
        }
        else if(d[mid]<k) l=mid+1;
        else r=mid-1;
    }
    return res;
}
inline int findr(int l,int r,int k){
    int res=-1;
    while(l<=r){
        int mid=(l+r)/2;
        if(d[mid]==k){
            res=mid,l=mid+1;
        }
        else if(d[mid]<k) l=mid+1;
        else r=mid-1;
    }
    return res;
}
int calc(int u,int lenth,int k){
    dis[u]=lenth,ecnt=0;
    getdis(u,0);
    sort(d+1,d+ecnt+1);
    int sum=0;
    for(int i=1;i<=ecnt;++i){
        if(d[i]*2>k) break;
        int l=findl(i,ecnt,k-d[i]),r=findr(i,ecnt,k-d[i]);
        sum+=r-l+1;
    }
    return sum; 
}
int solve(int u){
    for(int i=1;i<=p;++i) ans[i]+=calc(u,0,q[i]);
    vis[u]=true;
    for(int e=adj[u];e;e=nxt[e]){
        int v=to[e];
        if(vis[v]) continue;
        for(int j=1;j<=p;++j) ans[j]-=calc(v,val[e],q[j]);
        maxn=siz[v];
        getroot(v,root=0);
        solve(root);
    }
}
int main(){ n=read(),p=read();
    for(int i=1;i<n;++i){
        int u=read(),v=read(),w=read();
        addedge(u,v,w);
    }
    for(int i=1;i<=p;++i) q[i]=read();
    maxn=son[0]=n;
    getroot(1,root=0);
    solve(root);
    for(int i=1;i<=p;++i){
        if(ans[i])cout<<"Yes"<<'\n';
        else cout<<"No"<<'\n';
    }
    return 0;   
}