给定一棵有 n 个点的树,询问树上距离为 k的点对是否存在。
我们可以每次选择树的中心,作为rt,将rt所有的儿子节点的子树拆开,用分治的思想去处理每个小子树的答案
代码
#include<bits/stdc++.h>
using namespace std;
const int maxn=10010,maxm=10001000;
int n,m,tmp[maxm],judge[maxm];
int sz[maxn],vis[maxn];
int head[maxn],que[maxn];
int size,maxp[maxn];
int tot,rt,dis[maxn];
int q[maxm],ynn[maxn];
struct edge
{
int next,to,dis;
}e[maxn<<1];
void addedge(int next,int to,int dis)
{
e[++size].to=to;
e[size].dis=dis;
e[size].next=head[next];
head[next]=size;
}
void getzx(int t,int fat)
{
int i,j;
sz[t]=1;
maxp[t]=0;
for(i=head[t];i;i=e[i].next)
{
j=e[i].to;
if(j==fat||vis[j]) continue;
getzx(j,t);
sz[t]+=sz[j];
maxp[t]=max(sz[j],maxp[t]);
}
maxp[t]=max(maxp[t],tot-sz[t]);
if(maxp[t]<maxp[rt]) rt=t;
}
void getdis(int t,int fat)
{
tmp[++tmp[0]]=dis[t];
int i,j,k;
for(i=head[t];i;i=e[i].next)
{
j=e[i].to;
k=e[i].dis;
if(vis[j]||j==fat) continue;
dis[j]=dis[t]+k;
getdis(j,t);
}
}
void calc(int t)
{
int p=0,i,j,k,l;
for(i=head[t];i;i=e[i].next)
{
j=e[i].to;
k=e[i].dis;
if(vis[j]) continue;
tmp[0]=0;
dis[j]=k;
getdis(j,t);
for(k=tmp[0];k;k--)
for(l=1;l<=m;l++) if(que[l]>=tmp[k]) ynn[l]|=judge[que[l]-tmp[k]];
for(k=tmp[0];k;k--) q[++p]=tmp[k],judge[tmp[k]]=1;
}
for(i=p;i;i--) judge[q[i]]=0;
}
void solve(int t)
{
int i,j;
vis[t]=judge[0]=1; calc(t);
for(i=head[t];i;i=e[i].next)
{
j=e[i].to;
if(vis[j]) continue;
tot=sz[j];
maxp[rt=0]=maxm;
getzx(j,0);
solve(rt);
}
}
int main()
{
int x,y,z;
scanf("%d%d",&n,&m);
for(int i=1;i<n;i++)
{
scanf("%d%d%d",&x,&y,&z);
addedge(x,y,z); addedge(y,x,z);
}
for(int i=1;i<=m;i++) scanf("%d",&que[i]);
maxp[rt=0]=n;
tot=n;
getzx(1,0);
solve(rt);
for(int i=1;i<=m;i++)
{
if(ynn[i]) printf("AYE\n");
else printf("NAY\n");
}
return 0;
}