小萌新刚学点分治200ms惨遭TLE

P3806 【模板】点分治 1

qzmoot @ 2024-07-11 16:15:06

#include <bits/stdc++.h>
#define pb push_back
using namespace std;
const int N=1e4+5,NN=1e7+5;
int n,m,siz[N],as[N],maxn[N],rt,idx,ye[NN],no[NN],dis[N],book[N],ans[N];
int st[NN],top;
int head[N],cnt;
struct node{
    int nxt,to,dis;
}a[N<<1];
int read()
{
    int x=0;
    char ch=getchar();
    while(ch<'0' || ch>'9')ch=getchar();
    while(ch>='0' && ch<='9'){x=x*10+ch-48;ch=getchar();}
    return x;
}
void add(int u,int v,int w)
{
    a[++cnt].to=v;
    a[cnt].dis=w;
    a[cnt].nxt=head[u];
    head[u]=cnt;
}
void zx(int u,int fa)
{
    siz[u]=1;
    maxn[u]=0;
    for(int i=head[u];i;i=a[i].nxt)
    {
        int v=a[i].to,w=a[i].dis;
        if(v==fa || book[v])
            continue;
        zx(v,u);
        siz[u]+=siz[v];
        maxn[u]=max(maxn[u],siz[v]);
    }
    maxn[u]=max(maxn[u],idx-siz[u]);
    if(maxn[u]<maxn[rt])
        rt=u;
}
void gdis(int u,int fa)
{
    ye[++ye[0]]=dis[u];
    for(int i=head[u];i;i=a[i].nxt)
    {
        int v=a[i].to,w=a[i].dis;
        if(v==fa || book[v])
            continue;
        dis[v]=dis[u]+w;
        gdis(v,u);
    }
}
void calc(int u)
{
    int res=0;
    for(int i=head[u];i;i=a[i].nxt)
    {
        int v=a[i].to,w=a[i].dis;
        if(book[v])
            continue;
        ye[0]=0;
        dis[v]=w;
        gdis(v,u);
        for(int j=ye[0];j;j--)
            for(int k=1;k<=m;k++)
                if(as[k]>=ye[j])
                    ans[k]|=no[as[k]-ye[j]];
        for(int j=ye[0];j;j--)
            st[++top]=ye[j],no[ye[j]]=1;
    }
    for(int i=top;i;i--)
        no[st[i]]=0;
}
void dfs(int u)
{
    book[u]=no[0]=1;
    calc(u);
    for(int i=head[u];i;i=a[i].nxt)
    {
        int v=a[i].to,w=a[i].dis;
        if(book[v])
            continue;
        idx=siz[v];
        rt=0;
        maxn[rt]=NN;
        zx(v,0);
        dfs(rt);
    }
}
int main()
{
    n=read(),m=read();
    for(int i=1;i<n;i++)
    {
        int u=read(),v=read(),w=read();
        add(u,v,w),add(v,u,w);
    }
    for(int i=1;i<=m;i++)
        as[i]=read();
    maxn[0]=n;
    zx(1,0); 
    dfs(rt);
    for(int i=1;i<=m;i++)
        if(ans[i])
            puts("AYE");
        else
            puts("NAY");
    return 0;
}

|