【模板】点分治

#include <bits/stdc++.h>
using namespace std;
#define INF 0x3f3f3f3f
const int N=1e+5;
int n,m,tot,sum,rt,ans;
struct node{int v,w,next;}e[N<<1];
int head[N],mson[N],size[N],dis[N],rem[N];
int vis[N],test[N],judge[N],q[N],query[N];
void add(int u,int v,int w)
{
    e[++tot].next=head[u];
    e[tot].v=v;
    e[tot].w=w;
    head[u]=tot;
}
void getrt(int u,int fa)
{
    size[u]=1; mson[u]=0;
    for(int i=head[u];i;i=e[i].next) 
    {
        int v=e[i].v;
        if(v==fa||vis[v]) continue;
        getrt(v,u);
        size[u]+=size[v];
        mson[u]=max(mson[u],size[v]);
    }
    mson[u]=max(mson[u],sum-size[u]);
    if(mson[u]<mson[rt]) rt=u;
}
void getdis(int u,int fa)
{
    rem[++rem[0]]=dis[u];
    for(int i=head[u];i;i=e[i].next)
    {
        int v=e[i].v;
        if(v==fa||vis[v])continue;
        dis[v]=dis[u]+e[i].w;
        getdis(v,u);
    }
}
void calc(int u)
{
    int p=0;
    for(int i=head[u];i;i=e[i].next)
    {
        int v=e[i].v;
        if(vis[v])continue;
        rem[0]=0; dis[v]=e[i].w;
        getdis(v,u);//处理u的每个子树的dis

        for(int j=rem[0];j;--j)//遍历当前子树的dis
        for(int k=1;k<=m;++k)//遍历每个询问
        if(query[k]>=rem[j])
        test[k]|=judge[query[k]-rem[j]];
        //如果query[k]-rem[j]d的路径存在就标记第k个询问

        for(int j=rem[0];j;--j)//保存出现过的dis于judge
        q[++p]=rem[j],judge[rem[j]]=1;
    }
    for(int i=1;i<=p;++i)//处理完这个子树就清空judge
    judge[q[i]]=0;//特别注意一定不要用memset，会T

}

void solve(int u)
{   
    //judge[i]表示到根距离为i的路径是否存在
    vis[u]=judge[0]=1;calc(u);//处理以u为根的子树
    for(int i=head[u];i;i=e[i].next)//对每个子树进行分治
    {
        int v=e[i].v;
        if(vis[v])continue;
        sum=size[v]; mson[rt=0]=INF;
        getrt(v,0); solve(rt);//在子树中找重心并递归处理
    }
}

int main()
{
    scanf("%d%d",&n,&m);
    for(int i=1,u,v,w;i<n;++i)
    {
        scanf("%d%d%d",&u,&v,&w);
        add(u,v,w);add(v,u,w);
    }
    for(int i=1;i<=m;++i) scanf("%d",&query[i]);
    mson[rt]=sum=n; //第一次先找整棵树的重心
    getrt(1,0); 
    solve(rt);//对树进行点分治
    for(int i=1;i<=m;++i)
    {
        if(test[i]) printf("AYE
");
        else printf("NAY
");
    }
    return 0;
}

#include <bits/stdc++.h> 
using namespace std;
typedef long long ll;
const int inf=1e9+10;
const int maxn=1e5+10;
int head[maxn],tot;
int root,allnode,ans,n,k;
int vis[maxn],deep[maxn],dis[maxn],siz[maxn],point[maxn];//deep[0]子节点个数(路径长度)，point为重心节点
struct node{
    int to,next,val;
}edge[maxn<<1];
void add(int u,int v,int w)//前向星存图
{
    edge[tot].to=v;
    edge[tot].next=head[u];
    edge[tot].val=w;
    head[u]=tot++;
}
void init()//初始化
{
    memset(head,-1,sizeof head);
    memset(vis,0,sizeof vis);
    tot=0;
}

void get_root(int u,int father)//重心
{
    siz[u]=1;point[u]=0;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==father||vis[v]) continue;
        get_root(v,u);//递归得到子树大小
        siz[u]+=siz[v];
        point[u]=max(point[u],siz[v]);//更新u节点的point
    }
    point[u]=max(point[u],allnode-siz[u]);//保存节点size
    if(point[u]<point[root]) root=u;//更新当前子树的重心
}

void get_dis(int u,int father)//获取子树所有节点与根的距离
{
    deep[++deep[0]]=dis[u];
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(v==father||vis[v]) continue;
        int w=edge[i].val;
        dis[v]=dis[u]+w;
        get_dis(v,u);
    }
}

int cal(int u,int now)
{
    dis[u]=now;deep[0]=0;
    get_dis(u,0);
    sort(deep+1,deep+deep[0]+1);
    int all=0;
    for(int l=1,r=deep[0];l<r;){//二分
        if(deep[l]+deep[r]<=k){
            all+=r-l;l++;
        }
        else r--;
    }
    return all;
}

void solve(int u)//以u为重心进行计算
{
    ans+=cal(u,0);//以当前u为重心的贡献
    vis[u]=1;
    for(int i=head[u];~i;i=edge[i].next){
        int v=edge[i].to;
        if(vis[v]) continue;
        ans-=cal(v,edge[i].val);//减去子树的影响
        allnode=siz[v];
        root=0;
        get_root(v,u);
        solve(root);
    }
}

int main()
{
    while(~scanf("%d%d",&n,&k)&&n&&k){
        init();
        for(int i=1;i<n;i++){
            int u,v,w;
            scanf("%d%d%d",&u,&v,&w);
            add(u,v,w);
            add(v,u,w);
        }
        root=ans=0;
        allnode=n;point[0]=inf;
        get_root(1,0);
        
        solve(root);
        printf("%d
",ans);
    }
    return 0;
}