点分治模板题

#include<bits/stdc++.h>
using namespace std;

const int maxn=2e4+10;
int head[maxn],ver[maxn],nxt[maxn],edge[maxn];
int tot;
int vis[maxn]; // 分治时用来标记哪个重心已经使用过 
int ans,size[maxn],root,sum; 
// 求重心时使用，ans记录目前标记重心最小子树大小，size记录子树大小 ，root为重心，sum为目前该分治的子树大小 
//map<int,int> m;
int m[100000000],mcnt,mm[maxn]; // m[i]用来标记每个子节点到根的距离为i是否存在,mm用来清空m时用不然会T。 
int dis[maxn],bian[maxn],cnt; // dis[i]记录子节点到目前根的距离,bian记录每个存在的距离 
int q[1005],judge[1005];  
int n,k;

void add(int x,int y,int z)
{
    ver[++tot]=y;
    edge[tot]=z;
    nxt[tot]=head[x];
    head[x]=tot;
}
void dfs_find(int x,int fa) // 求重心 
{
    size[x]=1;
    int max_part=0;
    for(int i=head[x]; i; i=nxt[i])
    {
        int y=ver[i];
        if(y==fa||vis[y]) continue;
        dfs_find(y,x);
        size[x]+=size[y];
        max_part=max(max_part,size[y]);
    }
    max_part=max(max_part,sum-size[x]);
    if(ans>max_part)
    {
        ans=max_part;
        root=x;
    }
}

void dfs(int x,int fa) 
{
    bian[++cnt]=dis[x];
    for(int i=head[x]; i; i=nxt[i])
    {
        int y=ver[i];
        if(vis[y]||y==fa) continue;
        dis[y]=edge[i]+dis[x];
        dfs(y,x);
    }
}
void solve(int x) // 分治为子树解决，每次给一个子树的重心。 
{
    mcnt=0;
    for(int i=head[x]; i; i=nxt[i])
    {
        cnt=0;
        int y=ver[i];
        dis[y]=edge[i];
        if(vis[y]) continue;
        dfs(y,x);
        for(int j=1; j<=cnt; j++)
        {
            for(int l=1; l<=k; l++)
                if(q[l]-bian[j]>=0&&(m[q[l]-bian[j]]==1||q[l]==bian[j]))
                    judge[l]=1;
        }
        for(int j=1; j<=cnt; j++)
        {
            m[bian[j]]=1;
            mm[++mcnt]=bian[j];
        }
    }
    for(int i=1;i<=mcnt;i++)
        m[mm[i]]=0;
}
void divide(int x) //  划分子树
{
    vis[x]=1;
    solve(x);
    for(int i=head[x]; i; i=nxt[i])
    {
        int y=ver[i];
        if(vis[y]) continue;
        ans=sum=size[y];                          
        dfs_find(y,0);
        divide(root);
    }
}

int main()
{
    scanf("%d%d",&n,&k);
    for(int i=1; i<=n-1; i++)
    {
        int a,b,c;
        scanf("%d%d%d",&a,&b,&c);
        add(a,b,c);
        add(b,a,c);
    }
    for(int i=1; i<=k; i++)
        scanf("%d",&q[i]);
    ans=sum=n;
    dfs_find(1,0);
    divide(root);
    for(int i=1; i<=k; i++)
    {
        if(judge[i]) printf("AYE\n");
        else printf("NAY\n");
    }
}