P3806 【模板】点分治1(题解)(点分治)

P3806 【模板】点分治1(题解)(点分治)

洛谷题目传送门

#include<iostream>
#include<cstdlib>
#include<cstdio>
#include<cmath>
#include<cstring>
#include<iomanip>
#include<algorithm>
#include<ctime>
#include<queue>
#include<stack>
#include<vector>
#define rg register
#define il inline
#define lst long long
#define ldb long double
#define N 10050
#define KK 10000050
using namespace std;
const int Inf=1e9;

int n,m,K,cnt,top;
int root,Max,tot;
struct EDGE{
    int to,nxt,v;
}ljl[N<<1];
int hd[N],ans;
int siz[N],vis[N];
int pot[KK],Q[N],dis[N];

il int read()
{
    rg int s=0,m=0;rg char ch=getchar();
    while(ch<'0'||ch>'9'){if(ch=='-')m=1;ch=getchar();}
    while(ch>='0'&&ch<='9')s=(s<<3)+(s<<1)+(ch^48),ch=getchar();
    return m?-s:s;
}

il void add(rg int p,rg int q,rg int o)
{
    ljl[++cnt]=(EDGE){q,hd[p],o};hd[p]=cnt;
}

void get_root(rg int now,rg int fm)
{
    siz[now]=1;rg int num=0;
    for(rg int i=hd[now];i;i=ljl[i].nxt)
    {
        rg int qw=ljl[i].to;
        if(qw==fm||vis[qw])continue;
        get_root(qw,now);
        siz[now]+=siz[qw];
        num=max(num,siz[qw]);
    }
    num=max(num,tot-siz[now]);
    if(Max>num)Max=num,root=now;
}

void get_dis(rg int now,rg int fm)
{
    Q[++top]=dis[now];
    for(rg int i=hd[now];i;i=ljl[i].nxt)
    {
        rg int qw=ljl[i].to;
        if(qw==fm||vis[qw])continue;
        dis[qw]=dis[now]+ljl[i].v;
        get_dis(qw,now);
    }
}

il void Query(rg int now,rg int base,rg int op)
{
    dis[now]=base,top=0;
    get_dis(now,0);
    for(rg int i=1;i<=top;++i)
        for(rg int j=i+1;j<=top;++j)
            pot[Q[i]+Q[j]]+=op;
}

void divide(rg int now)
{
    Query(now,0,1),vis[now]=1;
    rg int all=tot;
    for(rg int i=hd[now];i;i=ljl[i].nxt)
    {
        rg int qw=ljl[i].to;
        if(vis[qw])continue;
        Query(qw,ljl[i].v,-1);
        tot=siz[now]>siz[qw]?siz[qw]:all-siz[now];
        Max=Inf,root=0;
        get_root(qw,0),divide(root);
    }
}

int main()
{
    n=read(),m=read();
    for(rg int i=1;i<n;++i)
    {
        rg int p=read(),q=read(),o=read();
        add(p,q,o),add(q,p,o);
    }
    tot=n,Max=Inf;
    get_root(1,0);divide(root);
//	puts("YES");
    for(rg int i=1;i<=m;++i)
    {
        K=read();
        if(pot[K])puts("AYE");
        else puts("NAY");
    }
    return 0;
}