P3806 【模板】点分治1

https://www.luogu.org/problem/P3806

#include <iostream>
#include <cstdio>
#include <queue>
#include <algorithm>
#include <cmath>
#include <cstring>
#define inf 2147483647
#define N 1000010
#define p(a) putchar(a)
#define For(i,a,b) for(int i=a;i<=b;++i)
//by war
//2019.8.24
using namespace std;
int n,m,x,y,v,rt,sum,cnt;
int k[N],f[N],siz[N],d[N],ans[N];
bool vis[N],fg[N];
queue<int>q;
struct node{
    int n;
    int v;
    node *next;
}*e[N];

void in(int &x){
    int y=1;char c=getchar();x=0;
    while(c<'0'||c>'9'){if(c=='-')y=-1;c=getchar();}
    while(c<='9'&&c>='0'){ x=(x<<1)+(x<<3)+c-'0';c=getchar();}
    x*=y;
}
void o(int x){
    if(x<0){p('-');x=-x;}
    if(x>9)o(x/10);
    p(x%10+'0');
}

void push(int x,int y,int v){
    node *p;
    p=new node();
    p->n=y;
    p->v=v;
    if(e[x]==0)
        e[x]=p;
    else{
        p->next=e[x]->next;
        e[x]->next=p;
    }
}

void GET_ROOT(int x,int fa){
    siz[x]=1;f[x]=0;
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa||vis[i->n]) continue;
        GET_ROOT(i->n,x);
        f[x]=max(f[x],siz[i->n]);
        siz[x]+=siz[i->n];
    }
    f[x]=max(f[x],sum-siz[x]);
    if(f[x]<f[rt]) rt=x;
}

void GET_DIST(int x,int fa,int w){
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa||vis[i->n]) continue;
        d[++cnt]=w+i->v;
        GET_DIST(i->n,x,d[cnt]);
    }
}

void GET_ANS(int x,int fa){
    fg[0]=1;q.push(0);vis[x]=1;
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa||vis[i->n]) continue;
        d[++cnt]=i->v;
        GET_DIST(i->n,x,i->v);
        For(i,1,cnt)
            For(j,1,m)
                if(k[j]>=d[i]) ans[j]|=fg[k[j]-d[i]];
        For(i,1,cnt) q.push(d[i]),fg[d[i]]=1;
        cnt=0;
    }
    while(!q.empty())
        fg[q.front()]=0,q.pop();
}

void dfs(int x,int fa){
    GET_ANS(x,fa);
    for(node *i=e[x];i;i=i->next){
        if(i->n==fa||vis[i->n]) continue;
        sum=siz[i->n];
        rt=0;
        f[rt]=inf;
        GET_ROOT(i->n,x);
        GET_ROOT(rt,rt);
        dfs(rt,x);
    }
}

signed main(){
    in(n);in(m);
    For(i,1,n-1){
        in(x);in(y);in(v);
        push(x,y,v);
        push(y,x,v);
    }
    For(i,1,m)
        in(k[i]);
    f[rt]=inf;
    sum=n;
    GET_ROOT(1,1);
    GET_ROOT(rt,rt);
    dfs(rt,rt);
    For(i,1,m)
        if(ans[i])
            puts("AYE");
        else
            puts("NAY");
    return 0;
}