
题目链接:https://vjudge.net/problem/POJ-1741
题意:给出一棵树,求出树上距离不超过k的点对数量。
思路:点分治经典题。先找重心作为树根,然后求出子树中所有点到重心的距离dis[i],那么所有组合为dis[i]+dis[j]<=k,其中不合法组合为在重心的同一个子树内的情况,所以要减去在重心的子树中仍满足dis[i]+dis[j]<=k的情况。
AC代码:
#include<cstdio>
#include<algorithm>
#include<cstring>
using namespace std; const int inf=0x3f3f3f3f;
const int maxn=;
struct node1{
int v,w,nex;
}edge[maxn<<]; struct node2{
int x,y;
}arr[maxn]; int n,k,cnt,ans,head[maxn],sz[maxn],mson[maxn],Min,root,size;
int vis[maxn],t,tt,dis[maxn],pre[maxn]; void adde(int u,int v,int w){
edge[++cnt].v=v;
edge[cnt].w=w;
edge[cnt].nex=head[u];
head[u]=cnt;
} void getroot(int u,int fa){
sz[u]=,mson[u]=;
for(int i=head[u];i;i=edge[i].nex){
int v=edge[i].v;
if(vis[v]||v==fa) continue;
getroot(v,u);
sz[u]+=sz[v];
if(sz[v]>mson[u]) mson[u]=sz[v];
}
if(size-sz[u]>mson[u]) mson[u]=size-sz[u];
if(mson[u]<Min) Min=mson[u],root=u;
} void getdis(int u,int fa,int len){
dis[++t]=len;
for(int i=head[u];i;i=edge[i].nex){
int v=edge[i].v;
if(vis[v]||v==fa) continue;
getdis(v,u,len+edge[i].w);
}
} void solve(int x,int y,int f){
t=;
getdis(x,,y);
sort(dis+,dis+t+);
tt=,dis[]=-,pre[]=;
for(int i=;i<=t;++i)
if(dis[i]!=dis[i-]) arr[++tt].x=dis[i],arr[tt].y=,pre[tt]=pre[tt-]+;
else ++arr[tt].y,++pre[tt];
for(int i=;i<=tt&&arr[i].x<=k/;++i)
ans+=(arr[i].y-)*arr[i].y/*f;
for(int i=;i<=tt&&arr[i].x<k/;++i){
int l=i+,r=tt,mid;
while(l<=r){
mid=(l+r)>>;
if(arr[i].x+arr[mid].x<=k) l=mid+;
else r=mid-;
}
ans+=(arr[i].y)*(pre[r]-pre[i])*f;
}
} void fenzhi(int u,int ssize){
vis[u]=;
solve(u,,);
for(int i=head[u];i;i=edge[i].nex){
int v=edge[i].v;
if(vis[v]) continue;
solve(v,edge[i].w,-);
Min=inf,root=;
size=sz[v]<sz[u]?sz[v]:(ssize-sz[u]);
getroot(v,);
fenzhi(root,size);
}
} int main(){
while(scanf("%d%d",&n,&k),n||k){
cnt=ans=;
for(int i=;i<=n;++i)
head[i]=,vis[i]=;
for(int i=;i<n;++i){
int u,v,w;
scanf("%d%d%d",&u,&v,&w);
adde(u,v,w);
adde(v,u,w);
}
Min=inf,root=,size=n;
getroot(,);
fenzhi(root,n);
printf("%d\n",ans);
}
}