BZOJ_1316_树上的询问_点分治

#include <stdio.h>
#include <string.h>
#include <algorithm>
using namespace std;
#define N 10050
#define maxl 1000000
int head[N],to[N<<1],nxt[N<<1],val[N<<1],cnt,n,m;
int fag[N],siz[N],ask[150],tot,solved[150],a[N],b[N],dis[N],root;
bool used[N],buc[1000010][2];
inline void add(int u,int v,int w) {
    to[++cnt]=v; nxt[cnt]=head[u]; head[u]=cnt; val[cnt]=w; 
}
void getroot(int x,int y) {
    fag[x]=0; siz[x]=1;
    int i;
    for(i=head[x];i;i=nxt[i]) {
        if(to[i]!=y&&!used[to[i]]) {
            getroot(to[i],x);
            siz[x]+=siz[to[i]];
            fag[x]=max(fag[x],siz[to[i]]);
        }
    }
    fag[x]=max(fag[x],tot-siz[x]);
    if(fag[root]>fag[x]) {root=x;}
}
void getdep(int x,int y) {
    int i;
    if(dis[x]<=maxl){
        a[++a[0]]=dis[x];
        b[++b[0]]=dis[x];
    }
    siz[x]=1;
    for(i=head[x];i;i=nxt[i]) {
        if(to[i]!=y&&!used[to[i]]) {
            dis[to[i]]=dis[x]+val[i];
            getdep(to[i],x);
            siz[x]+=siz[to[i]]; 
        }
    }
}
void getsiz(int x,int y) {
    int i;
    siz[x]=1;
    for(i=head[x];i;i=nxt[i]) if(to[i]!=y&&!used[to[i]]) {
        getsiz(to[i],x);
        siz[x]+=siz[to[i]];
    }
}
void work(int x) {
    used[x]=1;
    int i,j,k;
    dis[x]=0;
    b[0]=0;
    b[++b[0]]=0;
    for(i=head[x];i;i=nxt[i]) if(!used[to[i]]) {
        a[0]=0;
        a[++a[0]]=0;
        dis[to[i]]=val[i];
        getdep(to[i],0);
        for(k=1;k<=m;k++) if(!solved[k]) {
            for(j=1;j<=a[0];j++) {
                if(a[j]<=ask[k]&&buc[ask[k]-a[j]][0])
                    solved[k]=1;  
            }
        }
        for(j=1;j<=a[0];j++) buc[a[j]][0]=1,buc[a[j]][1]=0;
    }
    for(i=1;i<=b[0];i++) {
        buc[b[i]][0]=buc[b[i]][1]=0;
    }
    getsiz(x,0);
    for(i=head[x];i;i=nxt[i]) if(!used[to[i]]) {
        tot=siz[to[i]];
        root=0;
        getroot(to[i],0);
        work(root); 
    }
}
int main() {
    scanf("%d%d",&n,&m);
    int i,x,y,z;
    for(i=1;i<n;i++) {
        scanf("%d%d%d",&x,&y,&z);
        add(x,y,z); add(y,x,z);
    }
    for(i=1;i<=m;i++) {
        scanf("%d",&ask[i]);
        if(!ask[i]) solved[i]=1;    
    }
    fag[0]=100000000;
    tot=n;
    getroot(1,0);
    work(root);
    for(i=1;i<=m;i++) {
        if(solved[i]) puts("Yes");
        else puts("No");    
    }
}