点分治

P3806 【模板】点分治 1

点分治是一种树上的分治算法，即每次找到当前子树的重心并将其作为子树的根，进而对该子树的部分问题进行递归求解，这样可以保证递归层数不超过\(\log {n}\)。

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

const int N=1e7+5;
int n,m,qr[105],dis[10005],mx[10005],rt,mmr[N*10],cnt,siz[10005],sum,has[N];
bool exi[N*10],vis[10005],ans[105];
struct edge{
    int v,w;
};
vector<edge> e[N];

void getmx(int u,int fa) //找重心，因为每一棵子树都要遍历一遍，所以以点的个数为划分依据
{
    siz[u]=1,mx[u]=0;
    for (edge nex:e[u])
    {
        if (nex.v==fa||vis[nex.v]) continue;
        getmx(nex.v,u);
        siz[u]+=siz[nex.v];
        mx[u]=max(mx[u],siz[nex.v]);
    }
    mx[u]=max(mx[u],sum-mx[u]);
    if (mx[u]<mx[rt]) rt=u; //使子树的大小尽可能小
}

void getsiz(int u,int fa)
{
    siz[u]=1;
    for (edge nex:e[u])
    {
        if (nex.v==fa||vis[nex.v]) continue;
        getsiz(nex.v,u);
        siz[u]+=siz[nex.v];
    }
}

void getdis(int u,int fa)
{
    has[++has[0]]=dis[u]; //用桶存当前子树中的点到重心的每一个可能的距离
    for (edge nex:e[u])
    {
        if (nex.v==fa||vis[nex.v]) continue;
        dis[nex.v]=dis[u]+nex.w;
        getdis(nex.v,u);
    }
}
inline void calc(int u)
{
    cnt=0; //如果用memset(exi,0,sizeof exi)，将会获得1e7*log(n)加上一坨超大常数的时间复杂度
    for (edge nex:e[u])
    {
        int v=nex.v,w=nex.w;
        if (vis[v]) continue;
        dis[v]=w;
        has[0]=0;
        getdis(v,u); //对每棵子树更新到重心的距离
        for (int i=1;i<=has[0];i++)
            for (int j=1;j<=m;j++) if (qr[j]>=has[i]) ans[j]|=exi[qr[j]-has[i]]; //将当前子树中点到重心的距离与之前子树点到重心的距离进行匹配
        for (int i=1;i<=has[0];i++) mmr[++cnt]=has[i],exi[has[i]]=1; //将“当前子树”更新为“之前子树”
    }
    for (int i=1;i<=cnt;i++) exi[mmr[i]]=0;
}
void solve(int u)
{
    vis[u]=1;
    getsiz(u,u);
    calc(u);
    for (edge nex:e[u]) //对每棵子树进行递归
    {
        if (vis[nex.v]) continue;
        mx[rt=0]=1e9;
        sum=siz[nex.v];
        getmx(nex.v,u),solve(rt);
    }
}

int main()
{
    ios::sync_with_stdio(false);
    cin.tie(0);
    cin>>n>>m;
    exi[0]=1;
    for (int i=1;i<n;i++)
    {
        int u,v,w;
        cin>>u>>v>>w;
        e[u].push_back({v,w});
        e[v].push_back({u,w});
    }
    for (int i=1;i<=m;i++) cin>>qr[i]; //离线操作，避免每一次都带有点分治的巨大常数
    mx[rt]=1e9,sum=n;
    getmx(1,1);solve(rt);
    for (int i=1;i<=m;i++) cout<<(ans[i]?"AYE":"NAY")<<'\n';
}