luogu3806 【模板】点分治1

#include <algorithm>
#include <iostream>
#include <cstring>
#include <cstdio>
using namespace std;
int n, m, uu, vv, ww, a[105], cnt, hea[10005], gra[10005], siz[10005];
int tot, sze, rot;
bool ans[105], vis[10005];
struct Edge{
    int too, nxt, val;
}edge[20005];
struct Node{
    int dis, fro;
}nd[10005];
bool cmp(Node x, Node y){
    return x.dis<y.dis;
}
void add_edge(int fro, int too, int val){
    edge[++cnt].nxt = hea[fro];
    edge[cnt].too = too;
    edge[cnt].val = val;
    hea[fro] = cnt;
}
void getRoot(int x, int fa){
    siz[x] = 1;
    gra[x] = 0;
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(t!=fa && !vis[t]){
            getRoot(t, x);
            siz[x] += siz[t];
            gra[x] = max(gra[x], siz[t]);
        }
    }
    gra[x] = max(gra[x], sze-siz[x]);
    if(gra[rot]>gra[x])	rot = x;
}
void dfs(int x, int fa, int fro, int dis){
    nd[++tot] = (Node){dis, fro};
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(t!=fa && !vis[t])
            dfs(t, x, fro, dis+edge[i].val);
    }
}
int bs(int x){
    int l=1, r=tot, re=0;
    while(l<=r){
        int mid=(l+r)>>1;
        if(nd[mid].dis<x)	l = mid + 1;
        else{
            re = mid;
            r = mid - 1;
        }
    }
    return re;
}
void getAns(int x){
    tot = 0;
    nd[++tot] = (Node){0, 0};
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(!vis[t])
            dfs(t, x, t, edge[i].val);
    }
    sort(nd+1, nd+1+tot, cmp);
    for(int i=1; i<=m; i++)
        if(!ans[i]){
            int l=1;
            while(l<tot && nd[l].dis+nd[tot].dis<a[i])	l++;
            while(l<tot && !ans[i]){
                if(a[i]-nd[l].dis<nd[l].dis)	break;
                int pos=bs(a[i]-nd[l].dis);
                while(pos<=tot && nd[pos].fro==nd[l].fro && nd[pos].dis+nd[l].dis==a[i])	pos++;
                if(nd[l].dis+nd[pos].dis==a[i])	ans[i] = true;
                l++;
            }
        }
}
void work(int x){
    vis[x] = true;
    getAns(x);
    for(int i=hea[x]; i; i=edge[i].nxt){
        int t=edge[i].too;
        if(!vis[t]){
            rot = 0;
            sze = siz[t];
            getRoot(t, 0);
            work(rot);
        }
    }
}
int main(){
    cin>>n>>m;
    for(int i=1; i<n; i++){
        scanf("%d %d %d", &uu, &vv, &ww);
        add_edge(uu, vv, ww);
        add_edge(vv, uu, ww);
    }
    for(int i=1; i<=m; i++)
        scanf("%d", &a[i]);
    gra[0] = 0x3f3f3f3f;
    sze = n;
    rot = 0;
    getRoot(1, 0);
    work(rot);
    for(int i=1; i<=m; i++)
        if(ans[i])	printf("AYE
");
        else	printf("NAY
");
    return 0;
}