利用树分治的原理,再加上扫描线的o(nm)更新,可以求出树上任意两点距离==k的点对个数
#include<iostream>
#include<cstdio>
#include<cstring>
#include<algorithm>
using namespace std;
#define N 20005
#define inf 1e9+7
int hou[N],xia[N],zhong[N],v[N],tot,max1,sz[N],i,rt,dis[N],q[N],n,m,w[N],x,y,z,ans[N];
bool vis[N];
void jia(int x,int y,int z)
{
++tot,hou[tot]=xia[x],zhong[tot]=y,xia[x]=tot,v[tot]=z;
//cout<<tot<<" "<<x<<" "<<y<<endl;
}
void jian(int x,int y,int z)
{
jia(x,y,z);
jia(y,x,z);
}
int zzx(int o,int fu)
{
sz[o]=1;
int maxx=0;
for(int i=xia[o];i!=-1;i=hou[i])
{
int nd=zhong[i];
if(nd==fu||vis[nd])continue;
zzx(nd,o);
sz[o]+=sz[nd];
maxx=max(maxx,sz[nd]);
}
maxx=max(maxx,tot-maxx);
if(maxx<max1){max1=maxx,rt=o;}
}
void dfs(int o,int fu,int d)
{
dis[o]=d;
q[++q[0]]=d;
for(int i=xia[o];i!=-1;i=hou[i])
{
int nd=zhong[i];
//cout<<o<<" "<<nd<<endl;
if(nd==fu||vis[nd])continue;
dfs(nd,o,d+v[i]);
}
}
void tj1(int o,int fu)
{
q[0]=0;
dfs(o,fu,0);
sort(q+1,q+1+q[0]);
for(int i=1;i<=m;i++)
{
int l=1;
int r=q[0];
while(l<r)
{
if(q[l]+q[r]==w[i]){ans[i]++;if(q[l]+q[r-1]==w[i])--r;else ++l;}
else if(q[l]+q[r]>w[i])--r;
else ++l;
}
}
}
void tj2(int o,int fu)
{
q[0]=0;
dfs(o,fu,dis[o]-dis[fu]);
sort(q+1,q+1+q[0]);
for(int i=1;i<=m;i++)
{
int l=1;
int r=q[0];
while(l<r)
{
if(q[l]+q[r]==w[i]){ans[i]--;if(q[l]+q[r-1]==w[i])--r;else ++l;}
else if(q[l]+q[r]>w[i])++l;
else --r;
}
}
}
void work(int o,int fu)
{
vis[o]=1;
tj1(o,fu);
for(int i=xia[o];i!=-1;i=hou[i])
{
int nd=zhong[i];
if(nd==fu)continue;
tj2(nd,o);
tot=sz[nd];
max1=inf;
zzx(nd,o);
work(rt,o);
}
}
int main()
{
memset(xia,-1,sizeof(xia));
scanf("%d%d",&n,&m);
for(i=1;i<n;i++)
{
scanf("%d%d%d",&x,&y,&z);
jian(x,y,z);
}
for(i=1;i<=m;i++)
{
scanf("%d",&w[i]);
}
tot=n;
max1=inf;
zzx(1,0);
work(rt,0);
for(i=1;i<=m;i++)
{
printf("%d ",ans[i]);
}
}